# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_tv_detection.ipynb (unless otherwise specified).

__all__ = ['DetectronTrain', 'train_cli', 'evaluation', 'InferLogo', 'DetectronInference']

# Cell
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import os
import json
import copy
import cv2
import torch
import random
import numpy as np
from tqdm import tqdm
from glob import glob
import pandas as pd
from pycocotools.coco import COCO

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor, DefaultTrainer
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer, ColorMode
from detectron2.data import MetadataCatalog, DatasetCatalog, build_detection_test_loader
from detectron2.structures import BoxMode
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data.datasets import register_coco_instances

from fastscript import call_parse
from fastscript import *
from fastcore.transform import store_attr

from .cocoeval import COCOeval
from .utils import load_yaml

import matplotlib.pyplot as plt

# Cell
class DetectronTrain():
    def __init__(self, cfg, dataset_config: dict):
        store_attr(self, "cfg, dataset_config")
        self.cfg = cfg

        ## set dataset names
        self.cfg.DATASETS.TRAIN = tuple(self.dataset_config["TRAIN_NAME"])
        self.cfg.DATASETS.TEST = tuple(self.dataset_config["VAL_NAME"])
        # check cuda availability
        if not torch.cuda.is_available(): self.cfg.MODEL.DEVICE = 'cpu';
        os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) # create dir
        self.meta = {} if not self.dataset_config["KEYPOINT"] else {"keypoint_names": self.dataset_config["KEYPOINT_NAMES"], "keypoint_flip_map": self.dataset_config["KEYPOINT_FLIP_MAP"]}
        ## register data
        self.register_train_data()
        self.register_val_data()

        ## set trainer
        self.trainer = DefaultTrainer(self.cfg)

    def train(self):
        self.trainer.resume_or_load(resume=False)
        self.trainer.train()

    def evaluation(self, val_coco_name: str, output_dir: str):
        for dataset in self.cfg.DATASETS.TEST:
            metadata = MetadataCatalog.get(dataset)
            evaluator = COCOEvaluator(dataset, self.cfg, False, output_dir=self.cfg.OUTPUT_DIR)
            val_loader = build_detection_test_loader(self.cfg, dataset)
            print(f"evaluation metrics on {dataset}")
            print(inference_on_dataset(self.trainer.model, val_loader, evaluator))

    def register_train_data(self):
        for num, dataset_name in enumerate(self.dataset_config["TRAIN_NAME"]):
            register_coco_instances(dataset_name, self.meta, self.dataset_config["TRAIN_COCO_JSON_LOC"][num], self.dataset_config["TRAIN_IMG_DIR"][num])

    def register_val_data(self):
        for num, dataset_name in enumerate(self.dataset_config["VAL_NAME"]):
            register_coco_instances(dataset_name, self.meta, self.dataset_config["VAL_COCO_JSON_LOC"][num], self.dataset_config["VAL_IMG_DIR"][num])

# Cell
@call_parse
def train_cli(base_cfg: Param("base config file of the detectron2 model", type=str)='COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml',
          cfg: Param("config file", type=str)='../configs/mask_only_exp1.yml',
          data_cfg: Param("Config file for data", type=str)="../assets/brands.yml"):

    #train model
    dataset_config = load_yaml(data_cfg)
    config = get_cfg()
    # load base config weights
    config.merge_from_file(model_zoo.get_config_file(base_cfg))
    ## merge config file
    config.merge_from_file(cfg)
    model = DetectronTrain(config, dataset_config)
    model.train()
    #evaluate model
    for i in dataset_config["VAL_NAME"]:
          model.evaluation(i, config.OUTPUT_DIR)

# Cell

def evaluation(pred_coco_json: str, actual_coco_json: str, iou_type: str):
    """
    :param pred_coco_json: json file path with predicted results
    :param actual_coco_json: json file path with actual results
    :param iou_type: choose from ["segm","bbox","keypoints"] # keypoints not tested
    """
    coco_gt = COCO(actual_coco_json)
    with open(pred_coco_json) as f:
        coco_results = json.load(f)
    assert len(coco_results) > 0
    if iou_type == "segm":
        coco_results = copy.deepcopy(coco_results)
        # When evaluating mask AP, if the results contain bbox, cocoapi will
        # use the box area as the area of the instance, instead of the mask area.
        # This leads to a different definition of small/medium/large.
        # We remove the bbox field to let mask AP use mask area.
        for c in coco_results:
            c.pop("bbox", None)
    coco_dt = coco_gt.loadRes(coco_results)
    coco_eval = COCOeval(coco_gt, coco_dt, iou_type)
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    # print(coco_eval.stats) #contains all metrics in an array

# Cell
class InferLogo():
    def __init__(self, data_config: str):
        store_attr(self, "data_config")
        self.config = load_yaml(data_config)

        ROOT = self.config['ROOT']
        model_cfg = os.path.join(ROOT, self.config['CONFIGS']['logo_identification'])
        model_path = os.path.join(ROOT, self.config['WEIGHTS']['logo_identification'])

        self.cfg = get_cfg()
        self.cfg.merge_from_file(model_zoo.get_config_file(self.config['BASE_CONFIGS']['detection']))
        self.cfg.merge_from_file(model_cfg)

        if not torch.cuda.is_available(): self.cfg.MODEL.DEVICE = 'cpu';
        self.cfg.MODEL.WEIGHTS = model_path
        self.predictor = DefaultPredictor(self.cfg)

    def predict(self, img_loc: str, visualize: bool):
        im = cv2.imread(img_loc) if isinstance(img_loc, str) else img_loc
        outputs = self.predictor(im)["instances"].to("cpu")
        scores = outputs.scores.numpy().tolist()
        if len(scores)<1:
            return [],[]
        value, idx = max((value, idx) for (idx, value) in enumerate(scores))
        pred = outputs.pred_classes.numpy().tolist()[idx] # considering prediction with highest score
        if visualize:
            self.visualize(im, outputs)
        return self.config['LABELS'][pred], value

    def visualize(self, im, outputs: dict, figsize: tuple= (12, 5)):
        v = Visualizer(im[:, :, ::-1],scale=0.5)
        out = v.draw_instance_predictions(outputs)
        plt.figure(figsize=figsize)
        plt.imshow(out.get_image()[:, :, ::-1]); plt.show()

# Cell
class DetectronInference():
    def __init__(self, config_file: str):
        self.cfg = get_cfg()
        self.config_file = config_file
        self.cfg.merge_from_file(model_zoo.get_config_file(self.config_file))
        self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
        self.cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(self.config_file)
        if not torch.cuda.is_available(): self.cfg.MODEL.DEVICE = 'cpu';
        self.predictor = DefaultPredictor(self.cfg)

    def predict(self, img_loc: str):
        im = cv2.imread(img_loc)
        outputs = self.predictor(im)
        return outputs

    def process_single_class(self, outputs: dict, index: int=62):
        t = outputs["instances"].pred_classes
        index = (t == index).nonzero() # In coco index value of TV is 62.
        if index.size()[0] == 0:
            return None
        boxes = outputs["instances"].pred_boxes[index.squeeze(1)].tensor.numpy()
        masks = outputs["instances"].pred_masks[index.squeeze(1)].numpy()
        return boxes, masks

    def visualize(self, im, outputs: dict, figsize: tuple= (12, 5)):
        v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]), scale=1.2)
        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        plt.figure(figsize=figsize)
        plt.imshow(out.get_image()[:, :, ::-1]); plt.show()