app.py

import gradio as gr
import cv2
import numpy as np
import os
import yaml
import sys

# Add the scripts folder to Python path
sys.path.append(os.path.join(os.path.dirname(__file__), "scripts"))

from model_factory import get_model
from plot_bboxes_save_images_and_yolo_predictions import process_image_frame

# Load config and class names
labels_path = 'labels.txt'
with open(labels_path, 'r') as f:
    class_names = {i: line.strip() for i, line in enumerate(f) if line.strip()}

with open("config.yaml", "r") as file:
    config = yaml.safe_load(file)

confidence = config['inference']['confidence']
use_nms = config['inference']['use_nms']
nms_thresh = config['inference']['nms_thresh']
plot_dets = config['inference']['plot_dets']
save_preds = config['inference']['save_preds']
show = config['inference']['show']

# Load model
model_name = "rfdetr"
model_version = "small"
pretrain_weights = f"./checkpoint_best_total.pth"
model = get_model(model_name, model_version, pretrain_weights)


def predict(img, confidence_slider, iou_slider, use_nms_checkbox):
    """
    Run object detection on uploaded image and return annotated image.
    """
    # Convert PIL/numpy image to OpenCV BGR format
    img_bgr = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    # Inference
    detections = model.infer(img_rgb, confidence_slider, use_nms_checkbox, iou_slider)

    # Process + draw bounding boxes (in-memory, no save)
    annotated_img = process_image_frame(
        img_bgr.copy(), detections, class_names,
        save_path=None,  # don’t save to disk
        plot_dets=plot_dets, save_preds=save_preds, show=show
    )

    return cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)  # back to RGB for Gradio


# Gradio interface
title = "Animal Object Detection HUB"
description = "Run object detection on animal images using RF-DETR model."
examples = [
    ["images/Deer_0147.jpg", confidence, nms_thresh, True],
    ["images/Duck_0049.jpg", confidence, nms_thresh, True],
    ["images/Fox_0280.jpg", confidence, nms_thresh, True],
    ["images/Parrot_0062.jpg", confidence, nms_thresh, True],
    ["images/Rabbit_0099.jpg", confidence, nms_thresh, True],
    ["images/Raccoon_0241.jpg", confidence, nms_thresh, True],
    ["images/Squirrel_0056.jpg", confidence, nms_thresh, True],
    ["images/Tiger_0229.jpg", confidence, nms_thresh, True],
    ["images/Bear_0007.jpg", confidence, nms_thresh, True],
    ["images/Red panda_0044.jpg", confidence, nms_thresh, True]
]


gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Input", sources=["upload", "webcam", "clipboard"]),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=confidence, label="RF-DETR Confidence Threshold"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=nms_thresh, label="NMS Threshold"),
        gr.Checkbox(label="Use NMS", value=use_nms)
    ],
    outputs=gr.Image(type="pil", label="RF-DETR"),
    title=title,
    description=description,
    examples=examples
).launch(ssr_mode=False)