"""
building-segmentation
Proof of concept showing effectiveness of a fine tuned instance segmentation model for deteting buildings.
"""

from transformers import DetrFeatureExtractor, DetrForSegmentation
from PIL import Image
import gradio as gr
import numpy as np
import torch
import torchvision
import detectron2

import itertools
import seaborn as sns

cfg = get_cfg()

def segment_buildings(input_image, confidence):

    cfg.MODEL.WEIGHTS = "model_weights/chatswood_buildings_poc.pth"  
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set a custom testing threshold
    predictor = DefaultPredictor(cfg)

    outputs = predictor(im)

    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    output = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    output_image = output.get_image()[:, :, ::-1])
    
    return(output_image)

# gradio components -inputs
gr_image_input = gr.inputs.Image()
gr_slider_confidence = gr.inputs.Slider(0,1,.1,.7,
                                        label='Set confidence threshold % for masks')
# gradio outputs
gr_image_output = gr.outputs.Image() 

# Create user interface and launch
gr.Interface(predict_building_mask, 
                inputs = [gr_image_input,gr_slider_confidence],
                outputs = gr_image_output,
                title = 'Building Segmentation',
                description = "An instance segmentation webapp using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone").launch()