Upload surgvlp_linear_evaluation.py with huggingface_hub

surgvlp_linear_evaluation.py

@@ -0,0 +1,277 @@
+import torch.nn as nn
+import argparse
+import torch
+import clip
+from PIL import Image
+import sys
+sys.path.append('../../../')
+from codes.datasets import build_dataset
+from codes.models import build_algorithm
+from mmengine.config import Config
+from transformers import AutoTokenizer
+from baselines.utils import calc_accuracy, calc_f1
+import torchmetrics
+import numpy as np
+from torch.utils.data import ConcatDataset
+import torch.optim as optim
+
+def process_text(text):
+    tokenizer_clinical = AutoTokenizer.from_pretrained('/gpfswork/rech/okw/ukw13bv/mmsl/biobert_pretrain_output_all_notes_150000')
+    ixtoword = {v: k for k, v in tokenizer_clinical.get_vocab().items()}
+    if type(text) == str:
+        text = [text]
+
+    processed_text_tensors = []
+    for t in text:
+
+        text_tensors = tokenizer_clinical(
+            t,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=77,
+        )
+        text_tensors["sent"] = [
+            ixtoword[ix] for ix in text_tensors["input_ids"][0].tolist()
+        ]
+        processed_text_tensors.append(text_tensors)
+
+    caption_ids = torch.stack([x["input_ids"] for x in processed_text_tensors])
+    attention_mask = torch.stack(
+        [x["attention_mask"] for x in processed_text_tensors]
+    )
+    token_type_ids = torch.stack(
+        [x["token_type_ids"] for x in processed_text_tensors]
+    )
+
+    if len(text) == 1:
+        caption_ids = caption_ids.squeeze(0).cuda()
+        attention_mask = attention_mask.squeeze(0).cuda()#.to(device)
+        token_type_ids = token_type_ids.squeeze(0).cuda()
+    else:
+        caption_ids = caption_ids.squeeze().cuda()
+        attention_mask = attention_mask.squeeze().cuda()
+        token_type_ids = token_type_ids.squeeze().cuda()
+
+    cap_lens = []
+    for txt in text:
+        cap_lens.append(len([w for w in txt if not w.startswith("[")]))
+
+    return {
+        "input_ids": caption_ids,
+        "attention_mask": attention_mask,
+        "token_type_ids": token_type_ids,
+        "cap_lens": cap_lens,
+    }
+
+def test(classifier, test_loader, model, args):
+    class_prompt=args.class_prompt
+
+    model.eval()
+
+    with open(class_prompt) as f:
+        lines = f.readlines()
+    f.close()
+
+    class_texts = [i.replace('\n', '') for i in lines]
+    class_texts = process_text(class_texts)
+    text_features = model(None, class_texts, mode='text')['text_emb'].cuda()
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+
+
+    total_acc = []
+    total_f1_phase = []
+    total_f1_phase_class = []
+
+    with torch.no_grad():
+        for test_loader in test_loaders:
+            probs_list = []
+            label_list = []
+
+            for i, data in enumerate(test_loader): 
+                frames = data['video'].cuda() # (1, M, T, C, H, W)
+                # B, M, T, C, H, W = frames.shape
+                B, C, H, W = frames.shape
+
+                frames = frames.view(-1, C, H, W)
+                image_features = model(frames, None, mode='video')['img_emb'] # (B*M*T, D)
+
+                probs = classifier(image_features)
+
+                # probs = probs / probs.norm(dim=-1, keepdim=True)
+                # probs = probs @ text_features.to(dtype=torch.float32).T
+
+                probs = probs.softmax(dim=-1) # (1, classes)
+                labels = data['label'].cuda()
+
+                probs_list.append(probs)
+                label_list.append(labels)
+
+
+            #
+            probs_list = torch.cat(probs_list, 0)
+            labels = torch.cat(label_list, 0)
+            
+            acc = calc_accuracy(probs_list, labels)
+            print('accuracy: ', acc)
+            f1_class, f1_average = calc_f1(probs_list, labels)
+            print('f1 average: ', f1_average)
+            print('f1 classes: ', f1_class)  
+
+            total_acc.append(acc)
+            total_f1_phase.append(f1_average)
+        print('f1 phase video-wise average ', np.mean(np.asarray(total_f1_phase)))
+        print('Acc video-wise average ', np.mean(np.asarray(total_acc)))
+
+
+
+def linear_evaluation(
+    train_loader: torch.utils.data.DataLoader,
+    val_loader: torch.utils.data.DataLoader,
+    model: torch.nn.Module,
+    num_classes: int
+) -> torch.nn.Module:
+    # Freeze the pre-trained model's parameters
+    for param in model.parameters():
+        param.requires_grad = False
+    
+    class_prompt=args.class_prompt
+    with open(class_prompt) as f:
+        lines = f.readlines()
+    f.close()
+
+    class_texts = [i.replace('\n', '') for i in lines]
+    class_texts = process_text(class_texts)
+    text_features = model(None, class_texts, mode='text')['text_emb'].cuda()
+    text_features /= text_features.norm(dim=-1, keepdim=True).to(dtype=torch.float32)
+
+    # Create a linear classifier
+    classifier = nn.Linear(2048, num_classes).cuda()
+    criterion = nn.CrossEntropyLoss().cuda()
+    optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001, weight_decay=0.0005)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=40)
+
+
+    # Training loop
+    model.eval()  # Ensure the model is in evaluation mode
+    for epoch in range(25):
+        for batch in train_loader:
+            inputs = batch['video'].cuda()
+            labels = batch['label'].cuda()
+
+            # Forward pass through the pre-trained model to get features
+            with torch.no_grad():
+                features = model(inputs, None, mode='video')['img_emb'] # (B*M*T, D)
+
+            features = features.to(dtype=torch.float32)
+            # Forward pass through the classifier
+            outputs = classifier(features)
+
+            # outputs_feat = outputs_feat / outputs_feat.norm(dim=-1, keepdim=True)
+            # outputs = outputs_feat @ text_features.T
+
+            loss = criterion(outputs, labels)
+            print(loss)
+
+            # Backward and optimize
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+        
+        # scheduler.step()
+
+        # Validation can be added here if needed
+        # classifier = classifier.eval()
+        # test(classifier, test_loaders, model, args)
+        # classifier = classifier.train()
+        
+    return classifier  # Return the trained classifier
+
+def get_args(description='CLIP'):
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument('--class_prompt', default='../class_prompt.txt', type=str, help='prompt for categories')
+    parser.add_argument('--dataset_config', default='./config.py', type=str, help='dataset config')
+    parser.add_argument('--batch_size', default=1, type=int, help='batch for testing')
+    parser.add_argument('--num_class', default=12, type=int, help='class for classification')
+    parser.add_argument('--checkpoint', default='', type=str, help='Checkpoint to load')
+    args = parser.parse_args()
+    return args, parser
+
+import torch.distributed as dist
+if __name__ == "__main__":
+
+    args, _ = get_args()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    configs = Config.fromfile(args.dataset_config)['config']
+
+    model = build_algorithm(configs.model_config).cuda()
+
+    ###### load weights
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/epoch0917.pth.tar')['state_dict']
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_3/epoch0089.pth.tar')['state_dict']
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_best/epoch0200_archive.pth.tar')['state_dict']
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_3/epoch0111.pth.tar')['state_dict'] # Action+Phase
+
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_best_4/epoch0170.pth.tar')['state_dict']
+
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_test_4/epoch0500.pth.tar')['state_dict']
+
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_best_4_rewrite/epoch0250.pth.tar')['state_dict'] ### HecVL
+
+    # state_dict = torch.load('/gpfswork/rech/okw/ukw13bv/mmsl/configs/Hierarchy_SurgVLP_best_4_rewrite_spell_1/epoch0120.pth.tar')['state_dict'] ### NIPS
+
+    state_dict = torch.load(args.checkpoint)['state_dict']
+
+
+
+    new_dict = {}
+    for k, v in state_dict.items():
+        if 'module.' in k:
+            new_dict[k[7:].replace('visual.model.', 'backbone_img.model.').replace('text_module.model.', 'backbone_text.model.').replace('visual.global_embedder','backbone_img.global_embedder')] = v
+    # .replace('visual.model.', 'backbone_img.model.').replace('text_module.model.', 'backbone_text.model.').replace('visual.global_embedder','backbone_img.global_embedder') # for old version of model, convert keys
+    a, b = model.load_state_dict(new_dict, strict=True)
+
+    # print(1, a)
+    # print(2, b)
+
+    model.eval()
+
+    train_datasets = [build_dataset(c) for c in configs.train_config]
+    train_dataset = ConcatDataset(train_datasets)
+
+    val_datasets = [build_dataset(c) for c in configs.val_config]
+    val_dataset = ConcatDataset(val_datasets)
+
+    test_datasets = [build_dataset(c) for c in configs.test_config]
+    # 40 videos --> 40 datasets
+
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        drop_last=False,
+        num_workers=4
+    )
+
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=4
+    )
+
+    test_loaders = [torch.utils.data.DataLoader(
+        test_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=0
+    ) for test_dataset in test_datasets] # 40 dataloaders
+    print(args)
+
+    classifier = linear_evaluation(train_loader, val_loader, model, args.num_class)
+
+    test(classifier, test_loaders, model, args)
+