# Author: Alexander Riedel 
# License: Unlicensed
# Link: https://github.com/alexriedel1/detectron2-GradCAM

from plots.gradcam.gradcam import GradCAM, GradCamPlusPlus
import detectron2.data.transforms as T
import torch
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.detection_utils import read_image
from detectron2.modeling import build_model
from detectron2.data.datasets import register_coco_instances

class Detectron2GradCAM():
  """
      Attributes
    ----------
    config_file : str
        detectron2 model config file path
    cfg_list : list
        List of additional model configurations
    root_dir : str [optional]
        directory of coco.josn and dataset images for custom dataset registration
    custom_dataset : str [optional]
        Name of the custom dataset to register
    """
  def __init__(self, config_file, cfg_list, root_dir=None, custom_dataset=None):
      # load config from file
      cfg = get_cfg()
      cfg.merge_from_file(config_file)

      if custom_dataset:
          register_coco_instances(custom_dataset, {}, root_dir + "coco.json", root_dir)
          cfg.DATASETS.TRAIN = (custom_dataset,)
          MetadataCatalog.get(custom_dataset)
          DatasetCatalog.get(custom_dataset)

      if torch.cuda.is_available():
          cfg.MODEL.DEVICE = "cuda"
      else:
          cfg.MODEL.DEVICE = "cpu"

      cfg.merge_from_list(cfg_list)
      cfg.freeze()

      self.cfg =  cfg

  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)).requires_grad_(True)
      inputs = {"image": image, "height": height, "width": width}
      return inputs

  def get_cam(self, img, target_instance, layer_name, grad_cam_type="GradCAM"):
      """
      Calls the GradCAM++ instance

      Parameters
      ----------
      img : str
          Path to inference image
      target_instance : int
          The target instance index
      layer_name : str
          Convolutional layer to perform GradCAM on
      grad_cam_type : str
          GradCAM or GradCAM++ (for multiple instances of the same object, GradCAM++ can be favorable)

      Returns
      -------
      image_dict : dict
        {"image" : <image>, "cam" : <cam>, "output" : <output>, "label" : <label>}
        <image> original input image
        <cam> class activation map resized to original image shape
        <output> instances object generated by the model
        <label> label of the 
      cam_orig : numpy.ndarray
        unprocessed raw cam
      """
      model = build_model(self.cfg)
      checkpointer = DetectionCheckpointer(model)
      checkpointer.load(self.cfg.MODEL.WEIGHTS)

      image = img
      input_image_dict = self._get_input_dict(image)

      if grad_cam_type == "GradCAM":
        grad_cam = GradCAM(model, layer_name)

      elif grad_cam_type == "GradCAM++":
        grad_cam = GradCamPlusPlus(model, layer_name)
      
      else:
        raise ValueError('Grad CAM type not specified')

      with grad_cam as cam:
        cam, cam_orig, output = cam(input_image_dict, target_category=target_instance)
      
      image_dict = {}
      image_dict["image"] = image
      image_dict["cam"] = cam
      image_dict["output"] = output[0]
      image_dict["label"] = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]).thing_classes[output[0]["instances"].pred_classes[target_instance]]
      return image_dict, cam_orig