from gradio.outputs import Label
from icevision.all import *
import PIL
import torch
import gradio as gr
import os

# Load model
class_map = ClassMap(['selected_variant'])
backbone = faster_rcnn.backbones.resnet_fpn.resnet18(pretrained=True)
model = faster_rcnn.model(backbone=backbone, num_classes=len(class_map))
model.load_state_dict(torch.load('object_localization_full-ancestry.model.pth', map_location=torch.device('cpu')))

def predict(
    model, image, detection_threshold: float = 0.5, mask_threshold: float = 0.5
):
    infer_ds = Dataset.from_images([image])

    batch, samples = faster_rcnn.build_infer_batch(infer_ds)
    preds = faster_rcnn.predict(
        model=model,
        batch=batch,
        detection_threshold=detection_threshold
    )
    return samples[0]["img"], preds[0]
    
def show_preds(input_image, display_list, detection_threshold):
    display_label = ("Label" in display_list)
    display_bbox = ("BBox" in display_list)

    if detection_threshold==0: detection_threshold=0.5

    img, pred = predict(model=model, image=input_image, detection_threshold=detection_threshold)
    # print(pred)
    img = draw_pred(img=img, pred=pred, class_map=class_map, display_label=display_label, display_bbox=display_bbox)
    img = PIL.Image.fromarray(img)
    # print("Output Image: ", img.size, type(img))
    return img

# Populate examples in Gradio interface
examples = [
    ['1.jpg'],
    ['2.jpg'],
    ['3.jpg']
]

display_chkbox = gr.inputs.CheckboxGroup(["Label", "BBox"], label="Display")
detection_threshold_slider = gr.inputs.Slider(minimum=0, maximum=1, step=0.1, default=0.5, label="Detection Threshold")

outputs = gr.outputs.Image(type="pil")

gr_interface = gr.Interface(
    fn=show_preds,
    inputs=["image", display_chkbox,  detection_threshold_slider],
    outputs=outputs,
    title='Selection Scan - Object Detection',
    examples=examples)
gr_interface.launch(inline=False, share=False, debug=True)