detectron2monitor.py

try:
    import detectron2
except:
    import os 
    os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
import torch
from detectron2.utils.logger import setup_logger
setup_logger()

from detectron2.config import get_cfg
import detectron2.data.transforms as T
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.modeling import build_model
from detectron2.data.detection_utils import read_image
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog


import numpy as np
import cv2
import os
import time
import pickle
import gradio as gr
import tqdm
import matplotlib.pyplot as plt
import io
from PIL import Image
torch.manual_seed(0)
np.random.seed(0)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

from models.regnet import build_regnet_fpn_backbone
import models.metadata as metadata

from utils_clustering import *

from base_cam import EigenCAM
from pytorch_grad_cam.utils.model_targets import FasterRCNNBoxScoreTarget

fullName2ab_dict = {'PASCAL-VOC':"voc", 'BDD100K':"bdd", 'KITTI':"kitti", 'Speed signs':"speed", 'NuScenes':"nu"}
ab2FullName_dict = {'voc':"PASCAL-VOC", 'bdd':"BDD100K", 'kitti':"KITTI", 'speed':"Speed signs", 'nu':"NuScenes"}
class Detectron2Monitor():
    def __init__(self, id, backbone, confidence_threshold=0.05):
        self.id, self.label_list = self._get_label_list(id)
        self.backbone = backbone
        self.confidence_threshold = confidence_threshold
        self.cfg, self.device, self.model = self._get_model()
        self.label_dict = {i:label for i, label in enumerate(self.label_list)}
        self.eval_list = ["ID-voc-OOD-coco", "OOD-open", "voc-val"] if self.id == "voc" else ["ID-bdd-OOD-coco", "OOD-open", "voc-ood", f"{self.id}-val"]
        MetadataCatalog.get("custom_dataset").set(thing_classes=self.label_list)
    
    def _get_label_list(self, id):
        id = fullName2ab_dict[id]
        if id == 'voc':      
            label_list = metadata.VOC_THING_CLASSES
        elif id == 'bdd':
            label_list = metadata.BDD_THING_CLASSES
        elif id == 'kitti':
            label_list = metadata.KITTI_THING_CLASSES
        elif id == 'speed' or id == 'prescan':
            label_list = metadata.SPEED_THING_CLASSES
        else:
            label_list = metadata.NU_THING_CLASSES
        return id, label_list

    def _get_model(self):
        cfg = get_cfg()
        cfg.merge_from_file(f"models/configs/vanilla_{self.backbone}.yaml")
        cfg.MODEL.WEIGHTS = f"models/weights/model_final_{self.backbone}_{self.id}.pth" 
        cfg.MODEL.DEVICE='cpu'
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(self.label_list)
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.confidence_threshold
        model = build_model(cfg)
        model.eval()
        checkpointer = DetectionCheckpointer(model)
        checkpointer.load(cfg.MODEL.WEIGHTS)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = model.to(device)
        return cfg, device, model

    def _inference(self, model, inputs):
        with torch.no_grad():
            images = model.preprocess_image(inputs)  
            features = model.backbone(images.tensor)  
            proposals, _ = model.proposal_generator(images, features, None)  # RPN

            features_ = [features[f] for f in model.roi_heads.box_in_features]
            box_features = model.roi_heads.box_pooler(features_, [x.proposal_boxes for x in proposals])
            box_features = model.roi_heads.box_head(box_features)  # features of all 1k candidates
            predictions = model.roi_heads.box_predictor(box_features)
            pred_instances, pred_inds = model.roi_heads.box_predictor.inference(predictions, proposals)
            pred_instances = model.roi_heads.forward_with_given_boxes(features, pred_instances)

            # output boxes, masks, scores, etc
            pred_instances = model._postprocess(pred_instances, inputs, images.image_sizes)  # scale box to orig size
            # features of the proposed boxes
            feats = box_features[pred_inds].cpu().numpy()   
        return pred_instances, feats

    def _load_monitors(self, clustering_algo, nb_clusters, eps=5, min_samples=10):
        if clustering_algo == "dbscan":
            with open(f"monitors/{self.id}/{self.backbone}/{clustering_algo}/eps{eps}_min_samples{min_samples}.pkl", 'rb') as f:
                monitors_dict = pickle.load(f)  
        else:
            with open(f"monitors/{self.id}/{self.backbone}/{clustering_algo}/{nb_clusters}.pkl", 'rb') as f:
                monitors_dict = pickle.load(f)
        return monitors_dict
    
    def _evaluate(self, clustering_algo, nb_clusters, eps, min_samples):
        dataset_name = f"{self.id}-val"
        with open(f'val_feats/{self.id}/{self.backbone}/{dataset_name}_feats_tp_dict.pickle', 'rb') as f:
            feats_tp_dict = pickle.load(f)
        with open(f'val_feats/{self.id}/{self.backbone}/{dataset_name}_feats_fp_dict.pickle', 'rb') as f:
            feats_fp_dict = pickle.load(f)
        monitors_dict = self._load_monitors(clustering_algo, nb_clusters, eps, min_samples)
        # make verdicts on ID data
        data_tp = []
        data_fp = []
        accept_sum = {"tp": 0, "fp": 0}
        reject_sum = {"tp": 0, "fp": 0}
        for label in tqdm.tqdm(self.label_list, desc="Evaluation on ID data"): 
            if label in monitors_dict:    
                verdict = monitors_dict[label].make_verdicts(feats_tp_dict[label])
                data_tp.append([label, len(verdict), np.sum(verdict)/len(verdict)])
                accept_sum["tp"] += np.sum(verdict)
                reject_sum["tp"] += len(verdict) - np.sum(verdict)   
                verdict = monitors_dict[label].make_verdicts(feats_fp_dict[label])
                data_fp.append([label, len(verdict), (len(verdict)-np.sum(verdict))/len(verdict)])
                accept_sum["fp"] += np.sum(verdict)
                reject_sum["fp"] += len(verdict) - np.sum(verdict)
        TPR = round((accept_sum['tp'] / (reject_sum['tp'] + accept_sum['tp'])*100), 2)
        FPR =  round((accept_sum['fp'] / (reject_sum['fp'] + accept_sum['fp'])*100), 2)
        id_name = ab2FullName_dict[self.id]
        df_id = pd.DataFrame([[id_name, f"{TPR}%", f"{FPR}%"]], columns=["Dataset", "TPR", "FPR"])
        
        data_ood = []
        i = 0
        self.eval_list.remove(dataset_name)
        for dataset_name in tqdm.tqdm(self.eval_list, desc="Evaluation on OOD data"):
            accept_sum = {"tp": 0, "fp": 0}
            reject_sum = {"tp": 0, "fp": 0}
            with open(f'val_feats/{self.id}/{self.backbone}/{dataset_name}_feats_fp_dict.pickle', 'rb') as f:
                feats_fp_dict = pickle.load(f)
            for label in self.label_list:
                if label in monitors_dict:
                    verdict = monitors_dict[label].make_verdicts(feats_fp_dict[label])
                    accept_sum["fp"] += np.sum(verdict)
                    reject_sum["fp"] += len(verdict) - np.sum(verdict)
            FPR =  round((accept_sum['fp'] / (reject_sum['fp'] + accept_sum['fp'])*100), 2)
            data_ood.append([dataset_name, str(FPR)+"%"])
            i += 1
        # prepare dataframes
        df_ood = pd.DataFrame(data_ood, columns=["Dataset", "FPR"])
        df_ood["Dataset"] = ["COCO", "Open Images"] if self.id == "voc" else ["COCO", "Open Images", "VOC-OOD"]
        return df_id, df_ood

    def _postprocess_cam(self, raw_cam, img_width, img_height):
        cam_orig = np.sum(raw_cam, axis=0)  # [H,W]
        cam_orig = np.maximum(cam_orig, 0)  # ReLU
        cam_orig -= np.min(cam_orig)
        cam_orig /= np.max(cam_orig)
        cam = cv2.resize(cam_orig, (img_width, img_height))
        return cam

    def _fasterrcnn_reshape_transform(self, x):
        target_size = x['p6'].size()[-2 : ]
        activations = []
        for key, value in x.items():
            activations.append(torch.nn.functional.interpolate(torch.abs(value), target_size, mode='bilinear'))
        activations = torch.cat(activations, axis=1)
        return activations
    
    def _get_input_dict(self, original_image):
        height, width = original_image.shape[:2]
        transform_gen = T.ResizeShortestEdge(
        [self.cfg.INPUT.MIN_SIZE_TEST, self.cfg.INPUT.MIN_SIZE_TEST], self.cfg.INPUT.MAX_SIZE_TEST
        )
        image = transform_gen.get_transform(original_image).apply_image(original_image)
        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
        inputs = {"image": image, "height": height, "width": width}
        return inputs
    
    def get_output(self, monitors_dict, img):
        image = read_image(img, format="BGR")
        input_image_dict = [self._get_input_dict(image)]
        pred_instances, feats = self._inference(self.model, input_image_dict)
        detections = pred_instances[0]["instances"].to("cpu")
        cls_idxs = detections.pred_classes.detach().numpy()
        # get labels from class indices
        labels = [self.label_dict[i] for i in cls_idxs]
        # count values in labels, and return a dictionary
        labels_count_dict = dict((i, labels.count(i)) for i in labels)
        v = Visualizer(image[..., ::-1], MetadataCatalog.get("custom_dataset"), scale=1)
        v = v.draw_instance_predictions(detections)
        img_detection = v.get_image()
        df = pd.DataFrame(list(labels_count_dict.items()), columns=['Object', 'Count'])
        verdicts = []
        for label, feat in zip(labels, feats):
            verdict = monitors_dict[label].make_verdicts(feat[np.newaxis,:])[0]
            verdicts.append(verdict)
        detections_ood = detections[[i for i, x in enumerate(verdicts) if not x]]
        detections_ood.pred_classes = torch.tensor([5]*len(detections_ood.pred_classes))
        labels_ood = [label for label, verdict in zip(labels, verdicts) if not verdict]
        verdicts_ood = ["Rejected"]*len(labels_ood)
        df_verdict = pd.DataFrame(list(zip(labels_ood, verdicts_ood)), columns=['Object', 'Verdict'])
        v = Visualizer(image[..., ::-1], MetadataCatalog.get("custom_dataset"), scale=1)
        for box in detections_ood.pred_boxes.to('cpu'):
            v.draw_box(box)
            v.draw_text("OOD", tuple(box[:2].numpy()))
        v = v.get_output()
        img_ood = v.get_image()
        pred_bboxes = detections.pred_boxes.tensor.numpy().astype(np.int32)
        target_layers = [self.model.backbone]
        targets = [FasterRCNNBoxScoreTarget(labels=labels, bounding_boxes=pred_bboxes)]
        cam = EigenCAM(self.model,
                        target_layers, 
                        use_cuda=False,
                        reshape_transform=self._fasterrcnn_reshape_transform)
        grayscale_cam = cam(input_image_dict, targets)
        cam = self._postprocess_cam(grayscale_cam, input_image_dict[0]["width"], input_image_dict[0]["height"])
        plt.rcParams["figure.figsize"] = (30,10)
        plt.imshow(img_detection[..., ::-1], interpolation='none')
        plt.imshow(cam, cmap='jet', alpha=0.5)
        plt.axis("off")
        img_buff = io.BytesIO()
        plt.savefig(img_buff, format='png', bbox_inches='tight', pad_inches=0)
        img_cam = Image.open(img_buff)
        image_dict = {}
        image_dict["image"] = image
        image_dict["cam"] = img_cam
        image_dict["detection"] = img_detection
        image_dict["verdict"] = img_ood
        return image_dict, df, df_verdict