app.py

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",  # Ensure the repo is accessible
    filename="model.pt",
    local_dir="./models"
)
model = YOLO(hf_hub_download(**repo_config))

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

def predict(image):
    try:
        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 "\n".join(results)  # Return as a single string for easier display in Gradio
    
    except Exception as e:
        return f"Error: {str(e)}"  # Basic error handling

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