import os os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" import torch from torch.utils.data import DataLoader import torchvision from torchvision import transforms import van import lightning as L from lightning.pytorch import seed_everything import lightning.pytorch.callbacks as callbk from lightning.pytorch.loggers import TensorBoardLogger import Loaders import numpy as np import torchmetrics import defs import argparse from pathlib import Path from tqdm import tqdm import pandas as pd def setup_seed(seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) seed_everything(seed, workers=True) torch.backends.cudnn.deterministic = True class RARP_NVB_Classification_Head(torch.nn.Module): def __init__(self, in_features:int, out_features:int, layer:list=[], activation_fn:torch.nn.Module = torch.nn.ReLU(), *args, **kwargs): super().__init__(*args, **kwargs) self.activation = activation_fn if len (layer) == 0: self.head = torch.nn.Linear(in_features, out_features) else: temp_head = [] next_input = in_features for num in layer: temp_head.append(torch.nn.Linear(next_input, num)) temp_head.append(self.activation) temp_head.append(torch.nn.Dropout(0.4)) next_input = num temp_head[-1] = torch.nn.Dropout(0.2) temp_head.append(torch.nn.Linear(next_input, out_features)) self.head = torch.nn.Sequential(*temp_head) del temp_head def forward(self, x): return self.head(x) class RARP_VAN(L.LightningModule): def __init__( self, van_model:str = "", lr:float = 1e-4, clasiffier_layers = [], lambda_L1:float = 1.31E-04 ): super().__init__() self.save_hyperparameters(ignore=["van_model"]) self.train_acc = torchmetrics.Accuracy('binary') self.val_acc = torchmetrics.Accuracy('binary') self.test_acc = torchmetrics.Accuracy('binary') self.f1ScoreTest = torchmetrics.F1Score('binary') if len(van_model) == 0: self.van_encoder = van.van_b2(pretrained=True, num_classes=0) print("pre-train ImageNet") else: self.van_encoder = van.van_b2(pretrained=False, num_classes=0) self.van_encoder.load_state_dict(torch.load(van_model)) self.image_emb = 512 self.clasiffier = RARP_NVB_Classification_Head(self.image_emb, 1, clasiffier_layers, torch.nn.SiLU(True)) self.lossFN_clasiffier = torch.nn.BCEWithLogitsLoss() def _calc_L1(self, params): l1 = 0 for p in params: l1 += torch.sum(torch.abs(p)) return self.hparams.lambda_L1 * l1 def forward(self, data): data = data.float() img_features = self.van_encoder(data) pred = self.clasiffier(img_features) return pred def _shared_step(self, batch, val_step:bool=False): data, label = batch label = label.float() prediction = self(data) prediction = prediction.flatten() predicted_labels = torch.sigmoid(prediction) loss = self.lossFN_clasiffier(prediction, label) if not val_step: loss += self._calc_L1(self.clasiffier.parameters()) if self.hparams.lambda_L1 is not None else 0 return loss, label, predicted_labels def training_step(self, batch, batch_idx): loss, true_labels, pred_labels = self._shared_step(batch, False) self.log("train_loss", loss, on_epoch=True) self.train_acc.update(pred_labels, true_labels) self.log("train_acc", self.train_acc, on_epoch=True, on_step=False) return loss def validation_step(self, batch, batch_idx): loss, true_labels, pred_labels = self._shared_step(batch, True) self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True) self.val_acc.update(pred_labels, true_labels) self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True) def test_step(self, batch, batch_idx): _, true_labels, predicted_labels = self._shared_step(batch, True) self.test_acc.update(predicted_labels, true_labels) self.f1ScoreTest.update(predicted_labels, true_labels) self.log("test_acc", self.test_acc, on_epoch=True, on_step=False) self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False) def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr) return [optimizer] class RARP_ConvNext(RARP_VAN): def __init__(self, model = "", lr = 0.0001, clasiffier_layers=[], lambda_L1 = 0.000131): super().__init__("", lr, clasiffier_layers, lambda_L1) if len(model) == 0: self.van_encoder = torchvision.models.convnext_small(weights=torchvision.models.ConvNeXt_Small_Weights.DEFAULT) self.van_encoder.classifier[-1] = torch.nn.Identity() else: self.van_encoder = torchvision.models.convnext_small() self.van_encoder.classifier[-1] = torch.nn.Identity() self.van_encoder.load_state_dict(torch.load(model)) self.image_emb = 768 self.clasiffier = RARP_NVB_Classification_Head(self.image_emb, 1, clasiffier_layers, torch.nn.SiLU(True)) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--Phase", default="train", type=str, help="'train' or 'eval'") parser.add_argument("--Fold", type=int, default=0) parser.add_argument("--Workers", type=int, default=5) parser.add_argument("--Log_Name", type=str, default="logs_debug", help="the name of the directory of the log chkp") parser.add_argument("-p", "--Pre_train", type=str, default="RARP") parser.add_argument("-w", "--Weigth", type=str, default="") parser.add_argument("-lv", "--Log_version", type=int) parser.add_argument("-e", "--Encoder", type=str, default="VAN") args = parser.parse_args() Dataset = Loaders.RARP_DatasetCreator( "./DataSet_Ando_All_no20Crop", FoldSeed=505, createFile=True, SavePath="./DataSet_AndoAll20_crop", Fold=5, removeBlackBar=False, ) Dataset.mean, Dataset.std = ([30.38144216, 42.03988769, 97.8896116], [40.63141752, 44.26910074, 50.29294373]) Dataset.CreateFolds() setup_seed(2023) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") Fold = args.Fold InitResize=(512, 512) batchSize = 8 numWorkers = args.Workers MaxEpochs = 100 LogFileName = args.Log_Name rootFile = Dataset.CVS_File.parent.parent/f"fold_{Fold}" checkPtCallback = [ callbk.ModelCheckpoint(monitor='val_acc', filename="RARP-{epoch}", save_top_k=10, mode='max'), callbk.EarlyStopping(monitor="val_loss", mode="min", patience=5) ] traintransform = torch.nn.Sequential( transforms.Resize(InitResize, antialias=True, interpolation=transforms.InterpolationMode.BICUBIC), transforms.RandomRotation( degrees=(-15, 15), fill=5 ), transforms.RandomResizedCrop( 224, scale=(0.4, 1), antialias=True, interpolation=transforms.InterpolationMode.BICUBIC ), transforms.RandomApply([ transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0)) ], 0.3), transforms.Normalize(Dataset.mean, Dataset.std) ).to(device) valtransform = torch.nn.Sequential( transforms.Resize(256, antialias=True, interpolation=transforms.InterpolationMode.BICUBIC), transforms.CenterCrop(224), transforms.Normalize(Dataset.mean, Dataset.std) ).to(device) testtransform = torch.nn.Sequential( transforms.Resize(256, antialias=True, interpolation=transforms.InterpolationMode.BICUBIC), transforms.CenterCrop(224), transforms.Normalize(Dataset.mean, Dataset.std) ).to(device) trainDataset = torchvision.datasets.DatasetFolder( str (rootFile/"train"), loader = defs.load_file_tensor, extensions = "npy", transform = traintransform ) valDataset = torchvision.datasets.DatasetFolder( str (rootFile/"val"), loader = defs.load_file_tensor, extensions = "npy", transform = valtransform ) testDataset = torchvision.datasets.DatasetFolder( str (rootFile/"test"), loader=defs.load_file_tensor, extensions="npy", transform=testtransform ) Train_DataLoader = DataLoader( trainDataset, batch_size=batchSize, num_workers=numWorkers, shuffle=True, drop_last=True, pin_memory=True, persistent_workers=numWorkers>0 ) Val_DataLoader = DataLoader( valDataset, batch_size=batchSize, num_workers=numWorkers, shuffle=False, pin_memory=True, persistent_workers=numWorkers>0 ) Test_DataLoader = DataLoader( testDataset, batch_size=batchSize, num_workers=numWorkers, shuffle=False, pin_memory=True, persistent_workers=numWorkers>0 ) match(args.Phase): case "train": Model = None match (args.Encoder): case "VAN": Model = RARP_VAN(args.Weigth if args.Pre_train == "RARP" else "", clasiffier_layers=[], lr=1e-4, lambda_L1=None)#lambda_L1=2.22e-6 case "ConvNext": Model = RARP_ConvNext(args.Weigth if args.Pre_train == "RARP" else "", clasiffier_layers=[], lr=1e-4, lambda_L1=None) trainer = L.Trainer( deterministic=True, accelerator='gpu', devices=1, logger=TensorBoardLogger(save_dir=f"./{LogFileName}"), log_every_n_steps=5, callbacks=checkPtCallback, max_epochs=MaxEpochs, ) print("Train Phase") trainer.fit(Model, train_dataloaders=Train_DataLoader, val_dataloaders=Val_DataLoader) trainer.test(Model, dataloaders=Test_DataLoader, ckpt_path="best") case "eval_all": print("Evaluation Phase") rows = [] pathCkptFile = Path(f"./{LogFileName}/lightning_logs/version_{args.Log_version}/checkpoints/") for ckpFile in pathCkptFile.glob("*.ckpt"): print(ckpFile.name) Model = None match (args.Encoder): case "VAN": Model = RARP_VAN.load_from_checkpoint(ckpFile) case "ConvNext": Model = RARP_ConvNext.load_from_checkpoint(ckpFile) Model.to(device) Model.eval() Predictions = [] Labels = [] with torch.no_grad(): for data, label in tqdm(iter(Test_DataLoader)): data = data.float().to(device) label = label.to(device) pred = Model(data).flatten() Predictions.append(torch.sigmoid(pred)) Labels.append(label) Predictions = torch.cat(Predictions) Labels = torch.cat(Labels) print(Predictions, Labels) acc = torchmetrics.Accuracy('binary').to(device)(Predictions, Labels) precision = torchmetrics.Precision('binary').to(device)(Predictions, Labels) recall = torchmetrics.Recall('binary').to(device)(Predictions, Labels) auc = torchmetrics.AUROC('binary').to(device)(Predictions, Labels) f1Score = torchmetrics.F1Score('binary').to(device)(Predictions, Labels) specificty = torchmetrics.Specificity("binary").to(device)(Predictions, Labels) table = [ ["0.5000", f"{acc.item():.4f}", f"{precision.item():.4f}", f"{recall.item():.4f}", f"{f1Score.item():.4f}", f"{auc.item():.4f}", f"{specificty.item():.4f}", ""] ] for i in range(2): aucCurve = torchmetrics.ROC("binary").to(device) fpr, tpr, thhols = aucCurve(Predictions, Labels) index = torch.argmax(tpr - fpr) th2 = (recall + specificty - 1).item() th2 = 0.5 if th2 <= 0 else th2 th1 = thhols[index].item() if i == 0 else th2 accY = torchmetrics.Accuracy('binary', threshold=th1).to(device)(Predictions, Labels) precisionY = torchmetrics.Precision('binary', threshold=th1).to(device)(Predictions, Labels) recallY = torchmetrics.Recall('binary', threshold=th1).to(device)(Predictions, Labels) specifictyY = torchmetrics.Specificity("binary", threshold=th1).to(device)(Predictions, Labels) f1ScoreY = torchmetrics.F1Score('binary', threshold=th1).to(device)(Predictions, Labels) #cm2 = torchmetrics.ConfusionMatrix('binary', threshold=th1).to(device) #cm2.update(Predictions, Labels) #_, ax = cm2.plot() #ax.set_title(f"NVB Classifier (th={th1:.4f})") table.append([ f"{th1:.4f}", f"{accY.item():.4f}", f"{precisionY.item():.4f}", f"{recallY.item():.4f}", f"{f1ScoreY.item():.4f}", f"{auc.item():.4f}", f"{specifictyY.item():.4f}", ckpFile.name ]) rows += table df = pd.DataFrame(rows, columns=["Youden", "Acc","Precision","Recall","F1","AUROC","Specificity","CheckPoint"]) df.style.highlight_max(color="red", axis=0) print(df)