import os os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1" import torch from torch.utils.data import DataLoader import torch.nn as nn import torch.nn.functional as F import torchvision from torchvision import transforms from transformers import AutoTokenizer, AutoModel 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 Models as M from pathlib import Path import pandas as pd import Loaders import numpy as np import torchmetrics import defs import argparse from tqdm import tqdm from softadapt import LossWeightedSoftAdapt, NormalizedSoftAdapt LLM = "emilyalsentzer/Bio_ClinicalBERT" PROMPT = [( "Post-prostatectomy robotic laparoscopic view of the pelvic surgical bed in a {age}-year-old patient " "(BMI {BMI}, PSA {PSA} ng/mL) with clinical stage {cT}, pathologic stage {pT} and Gleason score {GS}; " "prostate size {prostate_size} mm; operating time was {surgery_time} min (console time {console_time} min); " "blood loss {blood_loss} mL. Neurovascular bundle preserved: {NVB}" )] 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_DatasetFolder_LLM(torchvision.datasets.DatasetFolder): def __init__(self, root: str, loader, Extra_Data: pd.DataFrame, tokenizer = None, extensions = None, transform = None, target_transform = None, is_valid_file = None ) -> None: super().__init__(root, loader, extensions, transform, target_transform, is_valid_file) self.Extra_Data = Extra_Data self.tokenizer = tokenizer def __getitem__(self, index: int): path, target = self.samples[index] name = Path(path).name Extra_data = self.Extra_Data[self.Extra_Data["raw_name"] == name].fillna('').to_dict("records")[0] prompt_text = PROMPT[0].format(**Extra_data) text_data = self.tokenizer(prompt_text, padding="max_length", truncation=True, max_length=128, return_tensors="pt") for k in text_data: text_data[k] = text_data[k].squeeze(0) sample = self.loader(path) if self.transform is not None: sample = self.transform(sample) if self.target_transform is not None: target = self.target_transform(target) return (sample, text_data), target class RARP_CLIP_loss(nn.Module): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.lossFN = torch.nn.CrossEntropyLoss() def forward(self, clip_logits:torch.Tensor): image_logits, text_logits = clip_logits label = torch.arange(image_logits.size(0), device=image_logits.device) loss = (self.lossFN(image_logits, label) + self.lossFN(text_logits, label)) / 2 return loss class RARP_CLIP(nn.Module): def __init__( self, text_output_feat_dim:int, img_output_feat_dim:int, embed_dim:int, latent_space_dim:int ): super().__init__() self.image_latent_space = M.RARP_NVB_MLP(img_output_feat_dim, latent_space_dim, n_layers=2) self.text_latent_space = M.RARP_NVB_MLP(text_output_feat_dim, latent_space_dim, n_layers=2) self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1/0.07))) def forward(self, image_emb, text_emb): x_img = self.image_latent_space(image_emb) x_text = self.text_latent_space(text_emb) #Normalize x_img = F.normalize(x_img, dim=-1) x_text = F.normalize(x_text, dim=-1) #Scaled cosine logits scale = self.logit_scale.exp() logits_img = scale * x_img @ x_text.t() logits_text = logits_img.t() return logits_img, logits_text class RARP_VAN_BERT(L.LightningModule): def __init__( self, bert_model_name:str = LLM, van_model:str = "", lr:float = 1e-4, latent_space_dim:int = 512, hiden_dim:int = 256, clasiffier_layers = [], softAdptAlgo:int = 0, softAdptBeta:float = 0.1, lambda_L1:float = 1.31E-04, ): super().__init__() self.save_hyperparameters(ignore=["bert_model_name", "van_model"]) self.loss_weights = [1, 1] self.softAdapt = NormalizedSoftAdapt(softAdptBeta) if softAdptAlgo == 1 else LossWeightedSoftAdapt(softAdptBeta) self.loss_history = { 0 : [], #'loss_Binary' 1 : [], #'loss_CLIP' } self.train_acc = torchmetrics.Accuracy('binary') self.val_acc = torchmetrics.Accuracy('binary') self.test_acc = torchmetrics.Accuracy('binary') self.f1ScoreTest = torchmetrics.F1Score('binary') self.bert_llm = AutoModel.from_pretrained(bert_model_name) self.text_emb = 768 for parms in self.bert_llm.parameters(): parms.requires_grad = False 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.clip = RARP_CLIP(self.text_emb, self.image_emb, hiden_dim, latent_space_dim) self.clasiffier = M.RARP_NVB_Classification_Head(self.image_emb, 1, clasiffier_layers, torch.nn.SiLU(True)) self.lossFN_clasiffier = torch.nn.BCEWithLogitsLoss() self.lossFN_CLIP = RARP_CLIP_loss() def forward(self, data): img_data, text_data = data img_data = img_data.float() ids_tokens = text_data["input_ids"] att_mask = text_data["attention_mask"] img_features = self.van_encoder(img_data) x_text = self.bert_llm( input_ids=ids_tokens, attention_mask=att_mask ) x_text = x_text.last_hidden_state[:,0] #[CLS] token logits_img, logits_text = self.clip(img_features, x_text) pred = self.clasiffier(img_features) return pred, (logits_img, logits_text) def _calc_L1(self, params): l1 = 0 for p in params: l1 += torch.sum(torch.abs(p)) return self.hparams.lambda_L1 * l1 def _calc_weights(self, log_weights:bool = True): self.loss_weights = self.softAdapt.get_component_weights( torch.tensor(self.loss_history[0][:-1] if len(self.loss_history[0]) % 2 == 0 else self.loss_history[0]), torch.tensor(self.loss_history[1][:-1] if len(self.loss_history[1]) % 2 == 0 else self.loss_history[1]), verbose=False ) if log_weights: self.log("W_loss_img", self.loss_weights[0], on_epoch=True, on_step=False) self.log("W_loss_CLIP", self.loss_weights[1], on_epoch=True, on_step=False) self.loss_history = { 0: [], 1: [], } def _shared_step(self, batch, val_step:bool = False): data, label = batch label = label.float() prediction, clip_logits = self(data) prediction = prediction.flatten() predicted_labels = torch.sigmoid(prediction) loss_list = [ self.lossFN_clasiffier(prediction, label), self.lossFN_CLIP(clip_logits) ] if not val_step: if self.hparams.lambda_L1 is not None: loss_list[0] += self._calc_L1(self.clasiffier.parameters()) for i, l in enumerate(loss_list): self.loss_history[i].append(l.item()) loss = 0 for l, w in zip(loss_list, self.loss_weights): loss += w * l return loss, label, predicted_labels, loss_list def on_train_epoch_start(self): if self.current_epoch % 2 == 0 and self.current_epoch != 0: self._calc_weights() def training_step(self, batch, batch_idx): loss, true_labels, pred_labels, losses = 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) self.log("train_loss_img", losses[0], on_epoch=True, on_step=False) self.log("train_loss_clip", losses[1], on_epoch=True, on_step=False) return loss def validation_step(self, batch, batch_idx): loss, true_labels, pred_labels, losses = 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) self.log("val_loss_img", losses[0], on_epoch=True, on_step=False) self.log("val_loss_clip", losses[1], on_epoch=True, on_step=False) 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.AdamW(self.parameters(), lr=self.hparams.lr) #, weight_decay=self.Lambda_L2 return [optimizer] 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("-lv", "--Log_version", type=int) 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") DumpCSV = pd.read_csv(Dataset.CVS_File) Extradata = pd.read_csv(Path("./DataSet_Ando_All_no20Crop/data_source_prompt.csv")) DumpCSV["raw_name"] = "Img0-" + DumpCSV["id"].astype(str) + ".npy" DumpCSV = DumpCSV.drop(columns=["id", "path", "mean_1", "mean_2", "mean_3", "std_1", "std_2", "std_3"]) NewData = pd.merge(Extradata, DumpCSV, on="name") Fold = args.Fold InitResize=(512, 512) batchSize = 16 numWorkers = 5 MaxEpochs = 100 LogFileName = "log_X22" 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) bert_tokenizer = AutoTokenizer.from_pretrained(LLM) trainDataset = RARP_DatasetFolder_LLM( str (rootFile/"train"), loader = defs.load_file_tensor, Extra_Data = NewData, tokenizer = bert_tokenizer, extensions = "npy", transform = traintransform ) valDataset = RARP_DatasetFolder_LLM( str (rootFile/"val"), loader = defs.load_file_tensor, Extra_Data = NewData, tokenizer = bert_tokenizer, extensions = "npy", transform = valtransform ) testDataset = RARP_DatasetFolder_LLM( str (rootFile/"test"), loader=defs.load_file_tensor, Extra_Data = NewData, tokenizer = bert_tokenizer, 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 = RARP_VAN_BERT(LLM, "van_b2_teacher_90_D2.pth", clasiffier_layers=[]) 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 = RARP_VAN_BERT.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) pred = pred.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)