import gradio as gr
from ultralytics import YOLO
from PIL import Image
from huggingface_hub import hf_hub_download

# Download model from Hugging Face
print("Downloading model from Hugging Face...")
model_path = hf_hub_download(
    repo_id="ajeetsraina/clinical-dental-pathology-detector",
    filename="models/dental_pathology_detector.pt"
)
print(f"Model downloaded to: {model_path}")

# Load model
model = YOLO(model_path)

def analyze_xray(image):
    if image is None:
        return None, "Please upload an X-ray image"
    
    results = model.predict(image, conf=0.25)
    annotated = results[0].plot()
    annotated_pil = Image.fromarray(annotated)
    detections = results[0].boxes
    
    if len(detections) == 0:
        report = "✅ No significant pathologies detected"
    else:
        class_counts = {}
        for box in detections:
            cls_name = model.names[int(box.cls[0])]
            if cls_name not in class_counts:
                class_counts[cls_name] = 0
            class_counts[cls_name] += 1
        
        report = f"**Total Findings:** {len(detections)}\n\n"
        for cls_name, count in class_counts.items():
            emoji = {"Cavity": "⚠️", "Fillings": "✅", "Implant": "🦷", "Impacted Tooth": "⚠️"}.get(cls_name, "ℹ️")
            report += f"{emoji} **{cls_name}**: {count} detected\n"
    
    return annotated_pil, report

demo = gr.Interface(
    fn=analyze_xray,
    inputs=gr.Image(type="pil", label="Upload Dental X-ray"),
    outputs=[
        gr.Image(type="pil", label="Analyzed X-ray"),
        gr.Markdown(label="Clinical Report")
    ],
    title="🦷 Clinical Dental Pathology Detector",
    description="Upload dental X-ray to detect cavities, fillings, implants, and impacted teeth. **98.9% mAP@50 accuracy**",
    article="Model: [ajeetsraina/clinical-dental-pathology-detector](https://huggingface.co/ajeetsraina/clinical-dental-pathology-detector)"
)

demo.launch()