import os 
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')

import gradio as gr
import logging
from detectron2.engine import DefaultPredictor
import cv2
from detectron2.config import get_cfg
from utils import add_bboxes

config_file="config.yaml"
cfg = get_cfg()
cfg.merge_from_file(config_file)
cfg.MODEL.DEVICE="cpu"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
# cfg.MODEL.WEIGHTS = "checkpoints_model_final_imagenet_40k_synthetic.pth"

def predict(
    model,
    img
):
    if model=="40k synthetic":
        weights = "checkpoints_model_final_imagenet_40k_synthetic.pth"
    elif model == "100k synthetic":
        weights = "checkpoints_model_final_imagenet_100k_synthetic.pth"
    else:
        weights = "checkpoints_model_final_imagenet_5k_synthetic.pth"
    cfg.MODEL.WEIGHTS=weights
    predictor = DefaultPredictor(cfg)
    im = cv2.imread(img.name)
    output = predictor(im)
    img = add_bboxes(im, output['instances'].pred_boxes, scores=output['instances'].scores)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img

title = "Indoor Pet Detection"
description = "This is an application trained with synthetic data from Unity Computer Vision. We trained a single class object detection model to recognize dogs using images of pets randomly posed and placed in our home interior environment. There are multiple pre-trained models trained with 5k, 40k and 100k synthetic data that you can choose for inference. <p> For more information please refer - <a target='_blank' href='https://github.com/Unity-Technologies/Indoor-Pet-Detection'> Github Project </a> </p>"
examples = [
    ["5k synthetic", 'example.jpg'],
    ["40k synthetic", 'example.jpg'],
    ["100k synthetic", 'example.jpg'],
    ["5k synthetic", 'example-2.jpg'],
    ["40k synthetic", 'example-2.jpg'],
    ["100k synthetic", 'example-2.jpg']
]


gr.Interface(predict, [gr.inputs.Dropdown(["5k synthetic", "40k synthetic", "100k synthetic"]), gr.inputs.Image(type="file")], outputs=gr.outputs.Image(type="pil"),enable_queue=True, title=title,
    description=description,
    # article=article,
    examples=examples).launch(debug=True)