!pip install pyyaml==5.1
!pip install torch==1.8.0+cu101 torchvision==0.9.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.8/index.html

!pip install ds2ai
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow

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
import requests
import os
import time
from ds2ai import DS2
import ast
import zipfile
import torch, torchvision
from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
from detectron2.data.datasets import register_coco_instances

ds2 = DS2(apptoken="----")
project = ds2.get_project(14883)
model = project.models[0]
file_url = None
if not os.path.exists("./data"):
    if project.filePath:
        file_url = project.filePath
    else:
        label_project = ds2.get_labelproject(project.labelproject)
        async_task = label_project.export(is_get_image=True)
        for i in range(0, 1000):
            time.sleep(10)
            async_task = ds2.get_asynctask(async_task.id)
            if async_task.status == 100:
                file_url = async_task.outputFilePath
                break

    if not file_url:
        raise (Exception("Please upload the training file."))

    file_name = file_url.split("/")[-1]
    response = requests.get(file_url)
    with open(file_name, 'wb') as output:
        output.write(response.content)

    os.makedirs("./data", exist_ok=True)
    os.makedirs("./models", exist_ok=True)

    with zipfile.ZipFile(file_name, 'r') as zf:
        zf.extractall(path="./data")
        zf.close()

print("project.id")
print(project.id)

configFile = f"data/coco.json"
fileRoute = f"data/"


cocoData = None
configFileValid = None

if os.path.exists("data/cocovalid.json"):
    configFileValid = f"data/cocovalid.json"

try:
    register_coco_instances(f"{model.id}", {}, configFile, fileRoute)
    if configFileValid:
       register_coco_instances(f"{model.id}_valid", {}, configFileValid, fileRoute)
except:
    pass


cfg = get_cfg()

cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")
cfg.DATASETS.TRAIN = (f"{model.id}",)
if configFileValid:
    cfg.DATASETS.TEST = (f"{model.id}_valid",)
else:
    cfg.DATASETS.TEST = ()

cfg.DATALOADER.NUM_WORKERS = 1
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.SOLVER.BASE_LR = 0.02
cfg.SOLVER.MAX_ITER = 300
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128

if project.yClass:
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(ast.literal_eval(project.yClass))
cfg.OUTPUT_DIR = f"./models/"

trainer = DefaultTrainer(cfg)

trainer.resume_or_load(resume=False)
trainer.train()
evaluator = COCOEvaluator(f"{model.id}", ("bbox", "segm"), False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, f"{model.id}",)
bbox = inference_on_dataset(trainer.model, val_loader, evaluator).get("bbox")
print(bbox)

learn_path = f"./models/model_final.pth"
files = {'uploadedModel': open(learn_path, 'rb')}
values = {'apptoken': '----', 'project': 14883, 'bbox': bbox}  # TODO
r = requests.post('https://api.ds2.ai/predictmodelfromcolab/', files=files, data=values)  # TODO


image_sample_file_path = None
for root, dirs, images in os.walk(f"./data"):
    if '__MACOSX' in root:
        continue
    for image in images:
        if not image.lower().endswith((".jpg", ".jpeg", ".png")):
            continue
        image_sample_file_path = f"{root}/{image}"
        break
    if image_sample_file_path:
        break

im = cv2.imread(image_sample_file_path)
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])
print("You can go back to the ds2.ai to check the output details.")