import os

import gradio as gr
import numpy as np
import cv2

from my_models import YOLOV5CLIPModel, YOLOV8CLIPModel


def annotated_image(
    image: np.ndarray, label: str, conf: float, bbox: list
) -> np.ndarray:
    line_thickness = max(1, int(0.005 * max(image.shape[:2])))
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    image = cv2.rectangle(
        image,
        (bbox[0], bbox[1]),
        (bbox[2], bbox[3]),
        (255, 0, 0),
        thickness=line_thickness,
    )
    image = cv2.putText(
        image,
        f"{label} {conf:.2f}",
        (bbox[0], max(bbox[1] - 2 * line_thickness, 0)),
        cv2.FONT_HERSHEY_SIMPLEX,
        thickness=max(line_thickness // 2, 1),
        lineType=cv2.LINE_AA,
        color=(0, 0, 0),
        fontScale=max(0.5, 0.1 * line_thickness),
    )
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    return image


def detect_mosquito(image):
    label, conf, bbox = YOLOV8CLIPModel().predict(image)
    return annotated_image(image, label, conf, bbox)


description = """# [Mosquito Alert Competition 2023](https://www.aicrowd.com/challenges/mosquitoalert-challenge-2023) - 7th Place Solution

Welcome to my Hugging Face Space showcasing the performance of our model. 

This competition focused on detecting and classifying various mosquito species. 

The target species were:
- **Aedes aegypti** - Species
- **Aedes albopictus** - Species
- **Anopheles** - Genus
- **Culex** - Genus (Species classification is challenging, so it is provided at the genus level)
- **Culiseta** - Genus
- **Aedes japonicus/Aedes koreicus** - Species complex (Differentiating between these two species is particularly challenging).

> ***Note:** Only one mosquito will be annotated even if there are multiple mosquitoes in the image.*

## Experiment Details

All the details regarding the experiments and source code for the models can be found in the [GitHub repository](https://github.com/HCA97/Mosquito-Classifiction/tree/main).
"""

iface = gr.Interface(
    fn=detect_mosquito,
    description=description,
    inputs=gr.Image(),
    outputs=gr.Image(),
    allow_flagging="never",
    examples=[os.path.join("examples", f) for f in os.listdir("examples")],
    cache_examples=True,
)
iface.launch()