app.py

import os
import cv2
import numpy as np
import gradio as gr
from PIL import Image
import supervision as sv
from collections import Counter
from ultralytics import YOLOv10
from huggingface_hub import hf_hub_download


def download_models(model_id):
    hf_hub_download("pankaj-munde/FScout_2.0", filename=f"{model_id}", local_dir=f"./", token=os.environ["HF_TOKEN"])
    return f"./{model_id}"

model_path = download_models("FScout_yolo_0.2.pt")
fscout_yolo_model = YOLOv10(model_path)
COLOR_PALETTE = sv.ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
# ROOT_PATH = "/home/stinpankajm/workspace/CropDr_Internal/media/"
# FSCOUT_YOLO = ""f"{ROOT_PATH}Models/YOLOv10/FScout_yolo_0.2.pt"
# fscout_yolo_model = YOLOv10(FSCOUT_YOLO)

def fScout_yolo_prediction(image, threshold=0.25):

    bounding_box_annotator = sv.BoundingBoxAnnotator(color=COLOR_PALETTE)
    label_annotator = sv.LabelAnnotator(color=COLOR_PALETTE, text_color=sv.Color.white(), text_scale=0.6, text_thickness=2)
    numpy_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # final_data = crop_dr_inspection(results["boxes"], image, crop_name)
    results = fscout_yolo_model(source=numpy_array, conf=threshold)[0]
    detections = sv.Detections.from_ultralytics(results)

    labels_data = [results.names[d] for d in detections.class_id]
    n_labels = [l[:2] for l in labels_data]
    annotated_image = bounding_box_annotator.annotate(scene=numpy_array, detections=detections)
    frame = label_annotator.annotate(scene=annotated_image, detections=detections, labels=n_labels)
    fimage = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

   
    new_res = []
    if detections.area.any():
        
        result_data = Counter(labels_data)
        if result_data:
            
            for k, v in dict(result_data).items():
                tmp = {}
                tmp["Insect"] = k
                tmp["Count"] = v
                new_res.append(tmp)

        return fimage, new_res

    else:
        return fimage, []
    

# Create the Gradio interface
app = gr.Interface(
    fn=fScout_yolo_prediction, 
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=["image", "json"],
    title="Object Detection",
    description="Upload an image to detect objects."
)

# Launch the Gradio app
app.launch()