import cv2
import khandy
import numpy as np
import gradio as gr
from PIL import Image
from modelscope import snapshot_download
from insectid import InsectDetector, InsectIdentifier

MODEL_DIR = snapshot_download("MuGemSt/insecta", cache_dir="./insectid/__pycache__")


def infer(filename: str):
    if not filename:
        None, "请上传图片 Please upload a picture"

    detector = InsectDetector()
    identifier = InsectIdentifier()
    image = khandy.imread(filename)
    if image is None:
        return None

    if max(image.shape[:2]) > 1280:
        image = khandy.resize_image_long(image, 1280)

    image_for_draw = image.copy()
    image_height, image_width = image.shape[:2]
    boxes, confs, classes = detector.detect(image)
    text = "未知"
    for box, _, _ in zip(boxes, confs, classes):
        box = box.astype(np.int32)
        box_width = box[2] - box[0] + 1
        box_height = box[3] - box[1] + 1
        if box_width < 30 or box_height < 30:
            continue

        cropped = khandy.crop_or_pad(image, box[0], box[1], box[2], box[3])
        results = identifier.identify(cropped)
        print(results[0])
        prob = results[0]["probability"]
        if prob >= 0.10:
            text = "{} {}: {:.2f}%".format(
                results[0]["chinese_name"],
                results[0]["latin_name"],
                100.0 * results[0]["probability"],
            )

        position = [box[0] + 2, box[1] - 20]
        position[0] = min(max(position[0], 0), image_width)
        position[1] = min(max(position[1], 0), image_height)
        cv2.rectangle(
            image_for_draw,
            (box[0], box[1]),
            (box[2], box[3]),
            (0, 255, 0),
            2,
        )
        image_for_draw = khandy.draw_text(
            image_for_draw,
            text,
            position,
            font=f"{MODEL_DIR}/simsun.ttc",
            font_size=15,
        )

    outxt = text.split(":")[0] if ":" in text else text
    return Image.fromarray(image_for_draw[:, :, ::-1], mode="RGB"), outxt


if __name__ == "__main__":
    iface = gr.Interface(
        fn=infer,
        inputs=gr.Image(label="上传昆虫照片 Upload insect picture", type="filepath"),
        outputs=[
            gr.Image(label="识别结果 Recognition result"),
            gr.Textbox(label="最可能的物种 Best match", show_copy_button=True),
        ],
        title="图像文件格式支持 PNG, JPG, JPEG 和 BMP, 且文件大小不超过 10M<br>Image file format support PNG, JPG, JPEG and BMP, and the file size does not exceed 10M.",
        examples=[
            f"{MODEL_DIR}/examples/butterfly.jpg",
            f"{MODEL_DIR}/examples/beetle.jpg",
        ],
        allow_flagging="never",
        cache_examples=False,
    )

    iface.launch()