import os 
import sys

os.system("pip install gdown")

os.system("pip install imutils")

os.system('pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.9/index.html')

os.system("pip install git+https://github.com/cocodataset/panopticapi.git")

import gradio as gr
# check pytorch installation: 
import detectron2
from detectron2.utils.logger import setup_logger

# import some common libraries
import numpy as np
import cv2
import torch

# 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, ColorMode
from detectron2.data import MetadataCatalog
from detectron2.projects.deeplab import add_deeplab_config
coco_metadata = MetadataCatalog.get("coco_2017_val_panoptic")

# import kMaXDeepLab project
from kmax_deeplab import add_kmax_deeplab_config

from PIL import Image
import imutils

cfg = get_cfg()
cfg.MODEL.DEVICE='cpu'
add_deeplab_config(cfg)
add_kmax_deeplab_config(cfg)
cfg.merge_from_file("configs/coco/panoptic-segmentation/kmax_convnext_large.yaml")
os.system("gdown 1b6rEnKw4PNTdqSdWpmb0P9dsvN0pkOiN")
cfg.MODEL.WEIGHTS = './kmax_convnext_large.pth'
cfg.MODEL.KMAX_DEEPLAB.TEST.SEMANTIC_ON = True
cfg.MODEL.KMAX_DEEPLAB.TEST.INSTANCE_ON = True
cfg.MODEL.KMAX_DEEPLAB.TEST.PANOPTIC_ON = True
predictor = DefaultPredictor(cfg)

os.system("wget https://i.imgur.com/Vj17K5z.jpg")

def inference(img):
    im = cv2.imread(img)
    im = imutils.resize(im, width=512)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
    panoptic_result = v.draw_panoptic_seg(outputs["panoptic_seg"][0].to("cpu"), outputs["panoptic_seg"][1]).get_image()
    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
    instance_result = v.draw_instance_predictions(outputs["instances"].to("cpu")).get_image()
    v = Visualizer(im[:, :, ::-1], coco_metadata, scale=1.2, instance_mode=ColorMode.IMAGE_BW)
    semantic_result = v.draw_sem_seg(outputs["sem_seg"].argmax(0).to("cpu")).get_image()    
    return Image.fromarray(np.uint8(panoptic_result)).convert('RGB'),Image.fromarray(np.uint8(instance_result)).convert('RGB'),Image.fromarray(np.uint8(semantic_result)).convert('RGB')
    

title = "kMaX-DeepLab"
description = "Gradio demo for kMaX-DeepLab. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.01527' target='_blank'>kMaX-DeepLab</a> | <a href='https://github.com/google-research/deeplab2' target='_blank'>Github Repo</a></p>"

examples = [['Vj17K5z.jpg']]

gr.Interface(inference, inputs=gr.inputs.Image(type="filepath"), outputs=[gr.outputs.Image(label="Panoptic segmentation",type="pil"),gr.outputs.Image(label="instance segmentation",type="pil"),gr.outputs.Image(label="semantic segmentation",type="pil")], title=title,
    description=description,
    article=article,
    examples=examples).launch(enable_queue=True,cache_examples=True)