# -*- coding: utf-8 -*-
"""Detectron2_customized.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1x0W2vHI73IqnDLfP-Juiy8FmPMxCgrvn

# Install detectron2
"""

import datetime
# from google.colab import drive

tstamp = datetime.datetime.now().strftime("%Y-%m-%d_%H%M")
drive.mount("/content/drive/")# , force_remount=True)

# !pip install pyyaml==5.1
# !pip install imgaug==0.2.6 -q
import torch
import torchvision 

print("DONE")

print("Installing detectron2. Please wait...")
# !pip install 'git+https://github.com/facebookresearch/detectron2.git' -q
print("... done.")

print(f"Torch: {torch.__version__}")
print(f"Torchvision: {torchvision.__version__}")

TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
CUDA_VERSION = torch.__version__.split("+")[-1]
# Install detectron2 that matches the above pytorch version
# See https://detectron2.readthedocs.io/tutorials/install.html for instructions
# !pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/$CUDA_VERSION/torch$TORCH_VERSION/index.html
# If there is not yet a detectron2 release that matches the given torch + CUDA version, you need to install a different pytorch.
import detectron2
print("DONE")
# exit(0)  # After installation, you may need to "restart runtime" in Colab. This line can also restart runtime

# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger

setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow
import logging
loglevel = logging.INFO
logformat = "[%(levelname)s]\t%(funcName)15s: %(message)s"
logging.basicConfig(level=loglevel, format=logformat)
# logging.disable()  # ------------------------------------------------------------------------ LOGGING

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

# !pip install --upgrade --force-reinstall bioblu -q

from bioblu.detectron import detectron
from bioblu.ds_manage import ds_convert

import bioblu 
print(f"bioblu version: {bioblu.__version__}")

"""# Convert dataset"""

yolo_ds_root_dir = "/content/drive/MyDrive/datasets/dataset_04_gnejna_with_duplicates"
#yolo_ds_root_dir = "/content/drive/MyDrive/datasets/dataset_05_mini_gnejna"

# coco_ds_root_dir = yolo_ds_root_dir + "_coco"
# coco_ds_root_dir = "/root/"

materials_dict = {0: "trash"}
ds_convert.cvt_yolo_to_detectron(yolo_ds_root_dir, materials_dict=materials_dict)
print("DONE")

# Extract image dict lists from jsons
fpath_json_train = os.path.join(yolo_ds_root_dir + "_detectron", "annotations/instances_detectron_train.json")
fpath_json_valid = os.path.join(yolo_ds_root_dir + "_detectron", "annotations/instances_detectron_val.json")
# fpath_json_test = os.path.join(coco_ds_root_dir, "test/test.json")
print(f"Train json: {fpath_json_train}")
print(f"Valid json: {fpath_json_valid}")

train_imgs = detectron.create_detectron_img_dict_list(fpath_json_train)
valid_imgs = detectron.create_detectron_img_dict_list(fpath_json_valid)
# test_imgs = detectron.create_detectron_img_dict_list(fpath_json_test)

classes = materials_dict # .values()  # not just values to prevent subscriptable error
print("DONE")

"""# Prepare for training"""

from detectron2.engine import DefaultTrainer

# # Monkeypatch:
# def hardsetpatch(self, key, val):
#     super().__setattr__(key, val)
# Metadata.hardset = hardsetpatch
# Metadata.hardset(key, val)

ds_name_train = "instances_detectron_train"
ds_name_val = "instances_detectron_val"

# Clear registered datasets
try:
  DatasetCatalog.remove(ds_name_train)
  DatasetCatalog.remove(ds_name_val)
except KeyError:
  print("No previously registered datasets.")
else:
  print("Cleared previously registered datasets.")


cfg = get_cfg()

print("Initializing weights and model")
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_C4_1x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_C4_1x.yaml")  # Let training initialize from model zoo

print("Registering datasets...")
DatasetCatalog.register(ds_name_train, lambda: detectron.create_detectron_img_dict_list(fpath_json_train))
DatasetCatalog.register(ds_name_val, lambda: detectron.create_detectron_img_dict_list(fpath_json_valid))
# DatasetCatalog.register("trash_test", lambda: detectron.create_detectron_img_dict_list(fpath_json_test))

MetadataCatalog.get("instances_detectron_train").set(thing_classes=["trash"])
MetadataCatalog.get("instances_detectron_val").set(thing_classes=["trash"])

# microcontroller_metadata = MetadataCatalog.get("trash_train")
metadata_train = MetadataCatalog.get("instances_detectron_train")
metadata_valid = MetadataCatalog.get("instances_detectron_val")


cfg.DATASETS.TRAIN = ("instances_detectron_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 4000    # 300 iterations seems good enough for the tutorial dataset; you will need to train longer for a practical dataset
cfg.SOLVER.STEPS = []        # do not decay learning rate
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
# # NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.
#
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
print(f"Output dir: {os.path.abspath(cfg.OUTPUT_DIR)}")
# Save config file
cfg_savepath = "/content/output/cfg.yaml"
with open(cfg_savepath, 'w') as f:
  f.write(cfg.dump())

"""## Visualize examples"""

skip_viz = True
import json
if not skip_viz:
  for d in random.sample(train_imgs, 3):
    print(json.dumps(d, indent=4))
    fname = d["file_name"]
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=metadata_train, scale=0.5)  # ::-1 turns RGB into BGR
    out = visualizer.draw_dataset_dict(d)
    cv2_imshow(out.get_image()[:, :, ::-1])  # ::-1 turns BGR into RGB

"""Environment info:"""

#!wget -nc -nv https://github.com/facebookresearch/detectron2/raw/main/detectron2/utils/collect_env.py && python collect_env.py

"""# Train"""

# cfg.MODEL.DEVICE='cpu'  # Only use when no gpu available.

train = False

if train:
  trainer = DefaultTrainer(cfg) 
  trainer.resume_or_load(resume=False)
  trainer.train()
else:
  pass

"""# Save results"""

import datetime
import shutil
results_src = "/content/output"
results_dst = "/content/drive/MyDrive/colab_outputs/" + tstamp

shutil.copytree(results_src, results_dst)

"""# Inference and evaluation
(from colab example)
"""

# Inference should use the config with parameters that are used in training
# cfg now already contains everything we've set previously. We changed it a little bit for inference:

pretrained_model = "/content/drive/MyDrive/colab_outputs/2022-04-30_1030/model_final.pth"
if not train:
  cfg.MODEL.WEIGHTS = pretrained_model
else:
  cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # path to the model we just trained


cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold
predictor = DefaultPredictor(cfg)

"""Randomly visualize a few prediction results:"""

from detectron2.utils.visualizer import ColorMode
import cv2
# dataset_dicts = get_balloon_dicts("balloon/val")
dataset_dicts = detectron.create_detectron_img_dict_list(fpath_json_valid)

print(dataset_dicts)



for d in random.sample(dataset_dicts, 3):    
    print(d["file_name"])
    im = cv2.imread(d["file_name"])
    print(type(im))
    #print(f"Img. {d['file_name']}: {im.shape}")
    outputs = predictor(im)  # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
    v = Visualizer(im[:, :, ::-1],
                   metadata=metadata_valid,  # From wayyy up when registering the dataset
                   scale=0.5, 
                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models
    )
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(out.get_image()[:, :, ::-1])

"""Evaluation based on AP:"""

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader

evaluator = COCOEvaluator(ds_name_val, output_dir="./output")
val_loader = build_detection_test_loader(cfg, ds_name_val)
print(inference_on_dataset(predictor.model, val_loader, evaluator))

"""# Evaluation 
(my own)
"""

# from detectron2.evaluation import evaluator, inference_on_dataset
# from detectron2.evaluation import DatasetEvaluator, COCOEvaluator

# class Counter(DatasetEvaluator):
#   def reset(self):
#     self.count = 0
#   def process(self, inputs, outputs):
#     for output in outputs:
#       self.count += len(output["instances"])
#   def evaluate(self):
#     # save self.count somewhere, or print it, or return it.
#     return {"count": self.count}

# def get_all_inputs_outputs():
#   for data in data_loader:
#     yield data, model(data)

# evaluator = COCOEvaluator(dataset_name="trash_train",
#                           tasks="bbox",
#                           distributed=True,
#                           output_dir=cfg.OUTPUT_DIR,
#                           max_dets_per_image=1000,
#                           use_fast_impl=False,
#                           allow_cached_coco=False)
# evaluator.reset()
# for inputs, outputs in get_all_inputs_outputs():
#   evaluator.process(inputs, outputs)
# eval_results = evaluator.evaluate()

# eval_results = inference_on_dataset(model, data_loader, DatasetEvaluators([COCOEvaluator(...), Counter()]))

shutil.copytree(results_src, results_dst)