import gradio as gr
import numpy as np
import config
from utils import *
from pytorch_grad_cam.utils.image import show_cam_on_image
from yolov3 import YOLOv3LightningModel

ex1 = [[f'examples/{i}.jpg'] for i in range(1,8)]
ex2 = [[f'examples/{i}.jpg'] for i in range(8,15)]
scaled_anchors = config.scaled_anchors

model = YOLOv3LightningModel()
model.load_state_dict(torch.load("yolov3.pth", map_location="cpu"), strict=False)
model.eval()

@torch.inference_mode()
def YoloV3_classifier(image,  thresh=0.5,iou_thresh=0.5):
    transformed_image = config.transforms(image=image)["image"].unsqueeze(0)
    output = model(transformed_image)
    
    bboxes = [[] for _ in range(1)]
    for i in range(3):
        batch_size, A, S, _, _ = output[i].shape
        anchor = scaled_anchors[i]
        boxes_scale_i = cells_to_bboxes(
            output[i], anchor, S=S, is_preds=True
        )
        for idx, (box) in enumerate(boxes_scale_i):
            bboxes[idx] += box

    nms_boxes = non_max_suppression(
        bboxes[0], iou_threshold=iou_thresh, threshold=thresh, box_format="midpoint",
    )
    plot_img = draw_bounding_boxes(image.copy(), nms_boxes, class_labels=config.PASCAL_CLASSES)

    return plot_img

window1 = gr.Interface(
    YoloV3_classifier,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Slider(0, 1, value=0.5, step=0.1, label="Threshold", info="Set Threshold value"),
        gr.Slider(0, 1, value=0.5, step=0.1, label="IOU Threshold", info="Set IOU Threshold value"),
    ],
    outputs=[
        gr.Image(label="YoloV3 Object Detection"),
    ],
    examples=ex1,
)


def visualize_gradCam(image, target_layer=-5, show_cam=True, transparency=0.5):
    if show_cam:
        cam = YoloCAM(model=model, target_layers=[model.layers[target_layer]], use_cuda=False)
        transformed_image = config.transforms(image=image)["image"].unsqueeze(0)
        grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]
        img = cv2.resize(image, (416, 416))
        img = np.float32(img) / 255
        cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True, image_weight=transparency)
    else:
        cam_image = image
        
    return cam_image

window2 = gr.Interface(
    visualize_gradCam,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Slider(-5, -2, value=-3, step=-1, label="Network Layer", info="GRAD-CAM Layer to visualize?"),
        gr.Checkbox(label="GradCAM", value=True, info="Visualize Class Activation Maps ??"), 
        gr.Slider(0, 1, value=0.5, step=0.1, label="Transparency", info="Set Transparency of GRAD-CAMs"), 
    ],
    outputs=[
        gr.Image(label="Grad-CAM Visualization"),
    ],
    examples=ex2,
)

app = gr.TabbedInterface([window1, window2], ["YOLO V3 Detection", "GradCAM Visualization"])
app.launch()