# Copyright (c) Facebook, Inc. and its affiliates.
# Modified by Bowen Cheng from: https://github.com/facebookresearch/detectron2/blob/master/demo/demo.py
import argparse
import glob
import multiprocessing as mp
import os
#os.environ["CUDA_VISIBLE_DEVICES"] = ""
try:
    import detectron2
except ModuleNotFoundError:
    os.system('pip install git+https://github.com/facebookresearch/detectron2.git')

try:
    import segment_anything
except ModuleNotFoundError:
    os.system('pip install git+https://github.com/facebookresearch/segment-anything.git')

# fmt: off
import sys
sys.path.insert(1, os.path.join(sys.path[0], '..'))
# fmt: on

import tempfile
import time
import warnings

import cv2
import numpy as np
import tqdm

from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.projects.deeplab import add_deeplab_config
from detectron2.utils.logger import setup_logger

from cat_seg import add_cat_seg_config
from demo.predictor import VisualizationDemo
import gradio as gr
import torch
from matplotlib.backends.backend_agg import FigureCanvasAgg as fc

# constants
WINDOW_NAME = "MaskFormer demo"


def setup_cfg(args):
    # load config from file and command-line arguments
    cfg = get_cfg()
    add_deeplab_config(cfg)
    add_cat_seg_config(cfg)
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    if torch.cuda.is_available():
        cfg.MODEL.DEVICE = "cuda"
    cfg.freeze()
    return cfg


def get_parser():
    parser = argparse.ArgumentParser(description="Detectron2 demo for builtin configs")
    parser.add_argument(
        "--config-file",
        default="configs/vitl_swinb_384.yaml",
        metavar="FILE",
        help="path to config file",
    )
    parser.add_argument(
        "--input",
        nargs="+",
        help="A list of space separated input images; "
        "or a single glob pattern such as 'directory/*.jpg'",
    )
    parser.add_argument(
        "--opts",
        help="Modify config options using the command-line 'KEY VALUE' pairs",
        default=(
        [
        "MODEL.WEIGHTS", "model_final_cls.pth", 
        "MODEL.SEM_SEG_HEAD.TRAIN_CLASS_JSON", "datasets/voc20.json",
        "MODEL.SEM_SEG_HEAD.TEST_CLASS_JSON", "datasets/voc20.json",
        "TEST.SLIDING_WINDOW", "True",
        "MODEL.SEM_SEG_HEAD.POOLING_SIZES", "[1,1]",
        "MODEL.PROMPT_ENSEMBLE_TYPE", "single",
        "MODEL.DEVICE", "cpu",
        ]),
        nargs=argparse.REMAINDER,
    )
    return parser

def save_masks(preds, text):
    preds = preds['sem_seg'].argmax(dim=0).cpu().numpy() # C H W
    for i, t in enumerate(text):
        dir = f"mask_{t}.png"
        mask = preds == i
        cv2.imwrite(dir, mask * 255)

def predict(image, text, model_type):
    #import pdb; pdb.set_trace()
    #use_sam = True #
    use_sam = model_type != "CAT-Seg"
    
    predictions, visualized_output = demo.run_on_image(image, text, use_sam)
    #save_masks(predictions, text.split(','))
    canvas = fc(visualized_output.fig)
    canvas.draw()
    out = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(canvas.get_width_height()[::-1] + (3,))

    return out[..., ::-1]

if __name__ == "__main__":
    args = get_parser().parse_args()
    cfg = setup_cfg(args)
    global demo
    demo = VisualizationDemo(cfg)

    iface = gr.Interface(
        fn=predict,
        inputs=[gr.Image(), gr.Textbox(placeholder='background, cat, person'), ], #gr.Radio(["CAT-Seg", "Segment Anycat"], value="CAT-Seg")],
        outputs="image",
        description="""## Segment Anything with CAT-Seg!
Welcome to the Segment Anything with CAT-Seg! 

In this demo, we combine state-of-the-art open-vocabulary semantic segmentation model, CAT-Seg with SAM(Segment Anything) for semantically labelling mask predictions from SAM.

Please note that this is an optimized version of the full model, and as such, its performance may be limited compared to the full model. 

Also, the demo might run on a CPU depending on the demand, so it may take a little time to process your image.

To get started, simply upload an image and a comma-separated list of categories, and let the model work its magic!""")
    iface.launch()