roi_predictor.py

import os
import cv2
import numpy as np
import supervision as sv
from huggingface_hub import hf_hub_download
from ultralytics import YOLOE

# Hugging Face space model details
HF_TOKEN = os.environ.get("HF_TOKEN")
REPO_ID = "imtk/knee-landmarks"

# Download the ROI model weights
print("Downloading model...")
roi_model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename="roi_model/weights/best.pt",
    token=HF_TOKEN
)
print(f"ROI Model downloaded to: {roi_model_path}")

roi_model = YOLOE(roi_model_path)
roi_model.fuse()

def predict_rois(image):
    if image is None:
        return None, [], {}
    
    # Convert Gradio RGB image to BGR for model prediction
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Predict ROIs
    results = roi_model.predict(img_bgr, conf=0.05)[0]
    detections = sv.Detections.from_ultralytics(results)
    
    rois = {}
    for xyxy, _, conf, class_id, _, class_name_dict in detections:
        shape_name = str(class_name_dict["class_name"])
        pt1 = (int(xyxy[0]), int(xyxy[1]))
        pt2 = (int(xyxy[2]), int(xyxy[3]))
        
        # Crop the region from original RGB image
        roi = image[pt1[1]:pt2[1], pt1[0]:pt2[0]]
        if roi.size == 0:
            continue
            
        area = abs(xyxy[0] - xyxy[2]) * abs(xyxy[1] - xyxy[3])
        
        # If class exists, keep the one with higher confidence
        if shape_name in rois:
            if rois[shape_name]["conf"] > conf:
                continue
                
        rois[shape_name] = {
            "minx": xyxy[0],
            "miny": xyxy[1],
            "maxx": xyxy[2],
            "maxy": xyxy[3],
            "box": (pt1, pt2),
            "img": roi,
            "conf": conf,
            "area": area,
        }
    
    # Draw annotations on the original image
    annotated_img = image.copy()
    gallery_items = []
    
    for name, data in rois.items():
        pt1, pt2 = data["box"]
        
        if name.endswith("L"):
            color = (255, 0, 0)  # Red
        else:
            color = (0, 0, 255)  # Blue
        
        # Draw box on annotated image
        cv2.rectangle(annotated_img, pt1, pt2, color, 3)
        # Put text label
        cv2.putText(
            annotated_img,
            f"{name} ({data['conf']:.2f})",
            (pt1[0], max(pt1[1] - 10, 20)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.8,
            color,
            2
        )
        
        # Add cropped ROI to gallery
        gallery_items.append((data["img"], name))

    return annotated_img, gallery_items, rois