from typing import Tuple
from ultralytics import YOLO
from ultralytics.engine.results import Boxes
from ultralytics.utils.plotting import Annotator

import gradio as gr

cell_detector = YOLO("./weights/yolo_uninfected_cells.pt")
yolo_detector = YOLO("./weights/yolo_infected_cells.pt")
redetr_detector = YOLO("./weights/redetr_infected_cells.pt")

models = {"Yolo V11": yolo_detector, "Real Time Detection Transformer": redetr_detector}
# classes = {"Yolo V11": [0], "Real Time Detection Transformer": [1]}


def inference(image, model, conf) -> Tuple[str, str, str]:
    bboxes = []
    labels = []
    healthy_cell_count = 0
    unhealthy_cell_count = 0
    cells_results = cell_detector.predict(image, conf=0.4)
    selected_model_results = models[model].predict(
        image, conf=conf
    )

    for cell_result in cells_results:
        boxes: Boxes = cell_result.boxes
        healthy_cells_bboxes = boxes.xyxy.tolist()
        healthy_cell_count += len(healthy_cells_bboxes)
        bboxes.extend(healthy_cells_bboxes)
        labels.extend(["healthy"] * healthy_cell_count)

    for res in selected_model_results:
        boxes: Boxes = res.boxes
        unhealthy_cells_bboxes = boxes.xyxy.tolist()
        unhealthy_cell_count += len(unhealthy_cells_bboxes)
        bboxes.extend(unhealthy_cells_bboxes)
        labels.extend(["unhealthy"] * unhealthy_cell_count)

    annotator = Annotator(image, font_size=5, line_width=1)

    for box, label in zip(bboxes, labels):
        annotator.box_label(box, label)

    img = annotator.result()
    return (img, healthy_cell_count, unhealthy_cell_count)


ifer = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(label="Input Image", type="numpy"),
        gr.Dropdown(
            choices=["Yolo V11", "Real Time Detection Transformer"], multiselect=False
        ),
        gr.Slider(minimum=0.01, maximum=1)
    ],
    outputs=[
        gr.Image(label="Output Image", type="numpy"),
        gr.Textbox(label="Healthy Cells Count"),
        gr.Textbox(label="Infected Cells Count"),
    ],
    title="Blood Cancer Cell Detection and Counting"
)

ifer.launch(share=True)