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

# from matplotlib.pyplot import axis
os.system('pip install altair')
import altair
import gradio as gr

import torch
import numpy as np

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

# Creación del modelo
cfg = get_cfg()

# Es el modelo utilizado por Ekimetrics
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")

# En caso de que no puedas usar GPU y debas usar la CPU, tienes que especificarlo de esta manera
if not torch.cuda.is_available():
    cfg.MODEL.DEVICE='cpu'
    

model = DefaultPredictor(cfg)


title = '<center><img src = "https://images.squarespace-cdn.com/content/v1/5573469fe4b0061829d437e6/1591631182400-7DJR03RV6ZOCN0TBPRD7/white-deloitte-logo1.jpg" width="130" height="20"></center><p>Detectron2 Image Detection</p>'
description = 'Implementación de Detectron2 en la detección de imágenes. Sube una imagen, dale a submit y espera unos segundos a ver el output de la imagen con los objetos detectados'
article = '<p>Conoce más en: <a href="https://www2.deloitte.com/es/es/pages/strategy-operations/solutions/analytics-and-cognitive.html">Visita Deloitte AI&Data</a></p><p>Desarrollado por Carlos y Lucía</p>'

def inference(image):
    print(image.height)

    height = image.height

    img = np.array(image.resize((640, 500)))
    outputs = model(img)

    v = Visualizer(img, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    
    return out.get_image()

css1 = "body {background-image: url(r'https://www2.deloitte.com/content/dam/Deloitte/in/Images/promo_images/in-deloitte-logo-1x1-noexp.png');}"
gr.Interface(
    inference, 
    [gr.inputs.Image(type="pil", label="Input")], 
    gr.outputs.Image(type="numpy", label="Output"),
    title=title,
    description=description,
    article=article,
    css = css1,
    examples=[]).launch()