import cv2
import onnxruntime
import numpy as np
import gradio as gr
from typing import List, Tuple

class YOLOv9:
    def __init__(self,
                 model_path: str,
                 classes: List[str],
                 original_size: Tuple[int, int] = (640, 640),
                 device: str = "CPU") -> None:
        self.model_path = model_path
        self.classes = classes
        self.device = device
        self.image_width, self.image_height = original_size
        self.create_session()

    def create_session(self) -> None:
        self.session = onnxruntime.InferenceSession(self.model_path, providers=['CPUExecutionProvider'])
        self.input_shape = self.session.get_inputs()[0].shape
        self.input_height, self.input_width = self.input_shape[2:4]
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))

    def preprocess(self, img: np.ndarray) -> np.ndarray:
        image_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(image_rgb, (self.input_width, self.input_height))
        input_image = resized / 255.0
        return input_image.transpose(2, 0, 1)[np.newaxis, :, :, :].astype(np.float32)

    def xywh2xyxy(self, x):
        y = np.copy(x)
        y[..., 0] = x[..., 0] - x[..., 2] / 2  # x_min
        y[..., 1] = x[..., 1] - x[..., 3] / 2  # y_min
        y[..., 2] = x[..., 0] + x[..., 2] / 2  # x_max
        y[..., 3] = x[..., 1] + x[..., 3] / 2  # y_max
        return y 

    def postprocess(self, outputs, score_threshold: float, iou_threshold: float):
        predictions = np.squeeze(outputs).T
        scores = np.max(predictions[:, 4:], axis=1)
        predictions = predictions[scores > score_threshold, :]
        scores = scores[scores > score_threshold]
        class_ids = np.argmax(predictions[:, 4:], axis=1)

        boxes = predictions[:, :4]
        boxes = self.xywh2xyxy(boxes)
        boxes[:, [0, 2]] *= self.image_width / self.input_width
        boxes[:, [1, 3]] *= self.image_height / self.input_height
        boxes = boxes.astype(np.int32)

        detections = []
        indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), score_threshold, iou_threshold)
        if len(indices) > 0:
            indices = indices.flatten()
            for i in indices:
                bbox = boxes[i]
                detections.append({
                    "class_index": class_ids[i],
                    "confidence": scores[i],
                    "box": bbox,
                    "class_name": self.classes[class_ids[i]]
                })
        return detections

    def detect(self, img: np.ndarray, score_threshold: float, iou_threshold: float) -> List:
        input_tensor = self.preprocess(img)
        outputs = self.session.run(None, {self.session.get_inputs()[0].name: input_tensor})[0]
        return self.postprocess(outputs, score_threshold, iou_threshold)

    def draw_detections(self, img, detections: List):
        h, w = img.shape[:2]
        font_scale = max(w, h) / 1000
        for detection in detections:
            x1, y1, x2, y2 = detection['box']
            class_id = detection['class_index']
            confidence = detection['confidence']
            color = self.color_palette[class_id]
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            label = f"{self.classes[class_id]}: {confidence:.2f}"
            text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 2)[0]
            text_w, text_h = text_size
            cv2.rectangle(img, (x1, y1 - text_h - 5), (x1 + text_w, y1), color, -1)
            cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 2)


def process_image(image: np.ndarray, score_threshold: float, iou_threshold: float) -> np.ndarray:
    if image is None:
        print("Tidak ada gambar yang diunggah.")
        return None
    
    weight_path = "best.onnx"
    classes = ["RBC", "WBC", "Difficult", "P-Falciparum", "P-Malariae", "P-Ovale", "P-Vivax"]
    
    detector = YOLOv9(model_path=weight_path, classes=classes, original_size=(image.shape[1], image.shape[0]))
    detections = detector.detect(image, score_threshold, iou_threshold)
    detector.draw_detections(image, detections)
    return image


# Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload an image"),
        gr.Slider(0.1, 1.0, value=0.1, step=0.05, label="Score Threshold"),
        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="IOU Threshold"),
    ],
    outputs=gr.Image(type="numpy", label="Detected Image"),
    live=True
)

iface.launch()