test.py

"""
Project: CVS-AdaptNet
-----
Copyright (c) University of Strasbourg, All Rights Reserved.
"""

import torch
import torch.nn as nn
import argparse
from torch import Tensor
import torchvision.models as models
from mmengine.config import Config
from mmengine.logging import MMLogger
from codes.datasets import build_dataset
from codes.models import build_algorithm
import json
import cv2
from tqdm import tqdm 
import numpy as np
from torchvision.transforms import transforms
from torch.utils.data import ConcatDataset
from torchmetrics import AveragePrecision as AP, Precision, Recall, F1Score
from codes.models.algorithms.CVS_AdaptNet import CVSAdaptNetModel, process_text
import matplotlib.pyplot as plt

def test_both_plot_8(test_loaders, model_1, args):
    class_prompt=args.class_prompt

    model_1.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_1(None, class_texts, mode='text')['text_emb'].cuda()
    text_features /= text_features.norm(dim=-1, keepdim=True)
        # init logger
    logger: MMLogger = MMLogger.get_current_instance()
    logger_info = []
    eval_results = {}
    probs_list = []
    label_list = []


    with torch.no_grad():
        for test_loader in test_loaders:
            for i, data in enumerate(test_loader): 
                probs_list_temp = []
                label_list_temp = []
                frames = data['video'].cuda() # (1, M, T, C, H, W)
                B, C, H, W = frames.shape

                frames = frames.view(-1, C, H, W)
                image_features = model_1(frames, None, mode='video')['img_emb'] # (B*M*T, D)
                image_features /= image_features.norm(dim=-1, keepdim=True)

                probs = (image_features @ text_features.T) # (1, classes)
        
                probabilities = probs.softmax(dim=-1) # (1, classes)


                # Define the indices for each class based on your earlier specification
                class1_indices = [4,5,6,7]  # Class 1 corresponds to these indices
                class2_indices = [2,3,6,7]  # Class 2 corresponds to these indices
                class3_indices = [1,3,5,7]  # Class 3 corresponds to these indices

                # Initialize an empty result tensor [1799, 3]
                result = torch.zeros((probabilities.shape[0], 3))

                # Sum the probabilities for each class
                result[:, 0] = torch.sum(probabilities[:, class1_indices], axis=1)  # Class 1
                result[:, 1] = torch.sum(probabilities[:, class2_indices], axis=1)  # Class 2
                result[:, 2] = torch.sum(probabilities[:, class3_indices], axis=1)  # Class 3

                gts = [[float(item) for item in json.loads(value)] for value in data['label']]

                probs_list = probs_list + [p.tolist() for p in result]
                #print(gts)

                label_list = label_list + [g for g in gts]


                

        torch_ap = AP(task='multilabel', num_labels=3, average='none')
        ds_preds = torch.stack([Tensor(p) for p in probs_list])
        ds_gt = torch.stack([Tensor(g).round() for g in label_list]).long()

        ds_ap = torch_ap(ds_preds, ds_gt)

        # log overall
        logger_info.append(f'ds_average_precision: {torch.nanmean(ds_ap):.4f}')
        eval_results['ds_average_precision'] = torch.nanmean(ds_ap)


        # log component-wise
        for ind, i in enumerate(ds_ap):
            logger_info.append(f'ds_average_precision_C{ind+1}: {i:.4f}')
            eval_results['ds_average_precision_C{}'.format(ind+1)] = i

        print('\n'.join(logger_info))
    return ds_ap


def test_both_plot_pos_neg(test_loaders, model_1, args):
    # Modify the string to create the negative class prompt path
    class_prompt = args.class_prompt
    class_prompt_neg = class_prompt.replace('.txt', '_no_class.txt')

    # class_prompt_neg will now have '_no_class' added before the '.txt' extension
    print(class_prompt_neg)
    model_1.eval()

    with open(class_prompt) as f:
        lines = f.readlines()
    f.close()

    with open(class_prompt_neg) as f:
        neg_lines = f.readlines()
    f.close()

    class_texts_ = [i.replace('\n', '') for i in lines]
    class_texts_pos = process_text(class_texts_)
    text_features = model_1(None, class_texts_pos, mode='text')['text_emb'].cuda()
    text_features /= text_features.norm(dim=-1, keepdim=True)

    class_texts_neg_ = [i.replace('\n', '') for i in neg_lines]
    class_texts_neg = process_text(class_texts_neg_)
    text_features_neg = model_1(None, class_texts_neg, mode='text')['text_emb'].cuda()
    text_features_neg /= text_features_neg.norm(dim=-1, keepdim=True)
        # init logger
    logger: MMLogger = MMLogger.get_current_instance()
    logger_info = []
    eval_results = {}

    probs_list = []
    label_list = []

    with torch.no_grad():
        for test_loader in test_loaders:
            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
                image_name = data['image_name']
                #print("image_name", image_name)
                frames = frames.view(-1, C, H, W)
                image_features = model_1(frames, None, mode='video')['img_emb'] # (B*M*T, D)

                image_features /= image_features.norm(dim=-1, keepdim=True)

                probs_pos = (image_features @ text_features.T) # (1, classes)
                probs_neg = (image_features @ text_features_neg.T) # (1, classes)
                # Stack the tensors along a new dimension to apply softmax
                stacked_probs = torch.stack([probs_pos, probs_neg], dim=0)  # Shape: (2, 1, classes)

                # Apply softmax along the first dimension (dim=0) to normalize between pos and neg
                softmax_probs = torch.nn.functional.softmax(stacked_probs, dim=0)

                # Extract the softmax values
                softmax_pos = softmax_probs[0]  # Softmax for probs_pos

                # Decision step: Choose probs_pos if its softmax value is greater, else use (1 - probs_pos)
                final_probs = softmax_pos

                ## for average precision calculation
                preds = final_probs.cpu()
        

                gts = [[float(item) for item in json.loads(value)] for value in data['label']]

                probs_list = probs_list + [p.tolist() for p in preds]
                label_list = label_list + [g for g in gts]


        torch_ap = AP(task='multilabel', num_labels=3, average='none')
        ds_preds = torch.stack([Tensor(p) for p in probs_list])
        #ds_gt_pred = torch.stack([Tensor(g) for g in gt_prob_list])
        ds_gt = torch.stack([Tensor(g).round() for g in label_list]).long()

        ds_ap = torch_ap(ds_preds, ds_gt)
        #print(ds_ap)

        # log overall
        logger_info.append(f'ds_average_precision: {torch.nanmean(ds_ap):.4f}')
        eval_results['ds_average_precision'] = torch.nanmean(ds_ap)


        # log component-wise
        for ind, i in enumerate(ds_ap):
            logger_info.append(f'ds_average_precision_C{ind+1}: {i:.4f}')
            eval_results['ds_average_precision_C{}'.format(ind+1)] = i

        print('\n'.join(logger_info))

    return ds_ap

def test_both(test_loaders, model_1, args):
    class_prompt=args.class_prompt

    model_1.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_1(None, class_texts, mode='text')['text_emb'].cuda()
    text_features /= text_features.norm(dim=-1, keepdim=True)
        # init logger
    logger: MMLogger = MMLogger.get_current_instance()
    logger_info = []
    eval_results = {}

    probs_list = []
    label_list = []

    with torch.no_grad():
        for test_loader in test_loaders:
            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_1(frames, None, mode='video')['img_emb'] # (B*M*T, D)
                # image_features = torch.mean(image_features, dim=0, keepdim=True) # (1, D)
                image_features /= image_features.norm(dim=-1, keepdim=True)
                probs = (image_features @ text_features.T) # (1, classes)

                ## for average precision calculation
                preds = probs.cpu()                
                gts = [[float(item) for item in json.loads(value)] for value in data['label']]
                #print(gts)
                probs_list = probs_list + [p.tolist() for p in preds]
                label_list = label_list + [g for g in gts]


        

        torch_ap = AP(task='multilabel', num_labels=3, average='none')
        ds_preds = torch.stack([Tensor(p) for p in probs_list]).sigmoid()
        ds_gt = torch.stack([Tensor(g).round() for g in label_list]).long()

        ds_ap = torch_ap(ds_preds, ds_gt)
        logger_info.append(f'ds_average_precision: {torch.nanmean(ds_ap):.4f}')
        eval_results['ds_average_precision'] = torch.nanmean(ds_ap)
        # log component-wise
        for ind, i in enumerate(ds_ap):
            logger_info.append(f'ds_average_precision_C{ind+1}: {i:.4f}')
            eval_results['ds_average_precision_C{}'.format(ind+1)] = i

        print('\n'.join(logger_info))
    return ds_ap

def get_args(description='CLIP'):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--class_prompt', default='../class_prompts/class_prompt_endoscapes_manual.txt', type=str, help='prompt for standard, positive_negative, multi_class')
    parser.add_argument('--checkpoint', default='../best_model.pth', type=str, help='prompt for categories')
    parser.add_argument('--dataset_config', default='../config/config_cvsadaptnet.py', type=str, help='dataset config')
    parser.add_argument('--batch_size', default=1, type=int, help='batch for testing')
    parser.add_argument('--inference_mode', default="standard", type=str, help='mode for inference')
    parser.add_argument('--save_dir', default="save", type=str, help='save directory')
    args = parser.parse_args()
    return args, parser

if __name__ == "__main__":

    args, _ = get_args()
    num_classes = 3

    configs = Config.fromfile(args.dataset_config)['config']
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    test_datasets = [build_dataset(c) for c in configs.test_config]
        
    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]

    model = build_algorithm(configs.model_config).cuda()
    state_dict = torch.load(args.checkpoint)['model_state_dict']

    new_dict = {}
    for k, v in state_dict.items():
        #if 'module.' in k:
        new_dict[k.replace('model.backbone_img.model.', 'backbone_img.model.').replace('model.backbone_text.model.', 'backbone_text.model.').replace('model.backbone_img.global_embedder','backbone_img.global_embedder')] = v
    a, b = model.load_state_dict(new_dict, strict=False)

    if args.inference_mode == "standard":
        test_both(test_loaders, model, args)
    elif args.inference_mode == "positive_negative":
        test_both_plot_pos_neg(test_loaders, model, args)
    elif args.inference_mode == "multi_class":
        test_both_plot_8(test_loaders, model, args)