from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import gradio as gr
import numpy as np
from PIL import Image
import pytesseract

# Load YOLO model from HuggingFace
repo_config = dict(
    repo_id="arnabdhar/YOLOv8-nano-aadhar-card",
    filename="model.pt",
    local_dir="./models"
)
model = YOLO(hf_hub_download(**repo_config))

# Get id-to-label mapping
id2label = model.names

def predict(image):
    image_np = np.array(image)
    detections = model.predict(image_np)[0]
    
    results = []
    for box, confidence, class_id in zip(detections.xyxy, detections.confidence, detections.class_id):
        x1, y1, x2, y2 = map(int, box)
        label = id2label[class_id]
        
        # Crop the detected region
        cropped_region = image_np[y1:y2, x1:x2]
        
        # Perform OCR on the cropped region
        ocr_text = pytesseract.image_to_string(cropped_region, config='--psm 6')
        
        results.append(f"Detected {label}: {ocr_text.strip()} with confidence {confidence:.2f}")
    
    return results

# Create Gradio interface
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text")
iface.launch()