app.py

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

MODEL_DIR = snapshot_download(
    "Genius-Society/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 = "Unknown"
    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]["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__":
    with gr.Blocks() as demo:
        gr.Interface(
            fn=infer,
            inputs=gr.Image(label="Upload insect picture", type="filepath"),
            outputs=[
                gr.Image(label="Recognition result", show_share_button=False),
                gr.Textbox(label="Best match", show_copy_button=True),
            ],
            title="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",
            ],
            flagging_mode="never",
            cache_examples=False,
        )

        gr.HTML(
            """
            <iframe src="//player.bilibili.com/player.html?bvid=BV14krgYJE4B&autoplay=0" scrolling="no" border="0" frameborder="no" framespacing="0" allowfullscreen="true" width="100%" style="aspect-ratio: 16 / 9;"></iframe>
            """
        )

    demo.launch()