app.py

"""
DENTO — Hugging Face Space
===========================
Tab 1: AlphaDent  — YOLOv11 nano+small ensemble with WBF
Tab 2: Dental Radiography — Faster R-CNN ResNet50 FPN

Weight repos:
  Elshawaf1/alphadent-weights   →  modelA.pt, modelB.pt
  Elshawaf1/alphadent-weights   →  fasterrcnn.pth

Upload fasterrcnn.pth from Kaggle once:
    api.upload_file(
        path_or_fileobj="/kaggle/working/model_epoch_9.pth",
        path_in_repo="fasterrcnn.pth",
        repo_id="Elshawaf1/alphadent-weights",
        repo_type="model"
    )
"""

import tempfile
import cv2
import gradio as gr
import numpy as np
import pandas as pd
import torch
import torchvision
from huggingface_hub import hf_hub_download
from PIL import Image
from ensemble_boxes import weighted_boxes_fusion
from torchvision import transforms
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from ultralytics import YOLO

# ── Shared config ──────────────────────────────────────────────────────────────
WEIGHTS_REPO = "Elshawaf1/alphadent-weights"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ── Tab 1: AlphaDent class names ───────────────────────────────────────────────
YOLO_CLASS_NAMES = {
    0: "Cavity",
    1: "Crack",
    2: "Calculus",
    3: "Missing Tooth",
    4: "Healthy",
}
YOLO_COLORS = [
    (60,  180, 255),
    (80,  255, 120),
    (255, 100,  80),
    (200,  60, 255),
    (60,  220, 180),
]

# ── Tab 2: Faster R-CNN class names ───────────────────────────────────────────
# Matches label_encoder.classes_ from the dental-radiography dataset
# Index 0 is background (required by Faster R-CNN)
RCNN_CLASS_NAMES = [
    "background",
    "Abscess",
    "Caries",
    "Crown",
    "Filling",
    "Implant",
    "Malaligned",
    "Mandibular Canal",
    "Missing teeth",
    "Periapical lesion",
    "Retained root",
    "Root Canal Treatment",
    "Root Piece",
    "Impacted tooth",
    "Maxillary sinus",
    "Bone Loss",
    "Fracture",
    "Permanent tooth",
    "Temporary restoration",
]
NUM_CLASSES_RCNN = len(RCNN_CLASS_NAMES)  # includes background

RCNN_COLORS = [
    (int(255 * (i / NUM_CLASSES_RCNN)),
     int(180 * ((i * 3) % NUM_CLASSES_RCNN) / NUM_CLASSES_RCNN),
     int(255 * (1 - i / NUM_CLASSES_RCNN)))
    for i in range(NUM_CLASSES_RCNN)
]

# ── Load models ────────────────────────────────────────────────────────────────
print("Downloading weights from HF Hub …")

# YOLO ensemble
try:
    path_a = hf_hub_download(repo_id=WEIGHTS_REPO, filename="modelA.pt")
    path_b = hf_hub_download(repo_id=WEIGHTS_REPO, filename="modelB.pt")
except Exception as e:
    print(f"YOLO weights not found ({e}), using pretrained.")
    path_a, path_b = "yolo11n.pt", "yolo11s.pt"

YOLO_CFGS = [
    {"model": YOLO(path_a), "imgsz": 640},
    {"model": YOLO(path_b), "imgsz": 800},
]

# Faster R-CNN
try:
    rcnn_path = hf_hub_download(repo_id=WEIGHTS_REPO, filename="fasterrcnn.pth")
    rcnn_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None)
    in_features = rcnn_model.roi_heads.box_predictor.cls_score.in_features
    rcnn_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES_RCNN)
    rcnn_model.load_state_dict(torch.load(rcnn_path, map_location=DEVICE))
    rcnn_model.to(DEVICE)
    rcnn_model.eval()
    RCNN_LOADED = True
    print("Faster R-CNN ready ✓")
except Exception as e:
    print(f"Faster R-CNN weights not found: {e}")
    RCNN_LOADED = False

print("Models ready ✓")


# ── Helper: draw boxes on image ────────────────────────────────────────────────
def draw_boxes(img_bgr, boxes_px, labels, scores, class_names, colors):
    rows = []
    for i, (box, label, score) in enumerate(zip(boxes_px, labels, scores)):
        cls_id = int(label)
        cls_name = class_names[cls_id] if cls_id < len(class_names) else f"class_{cls_id}"
        color = colors[cls_id % len(colors)]
        x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])

        cv2.rectangle(img_bgr, (x1, y1), (x2, y2), color, 2)
        label_text = f"{cls_name}  {score:.0%}"
        (tw, th), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
        pill_y1 = max(y1 - th - 8, 0)
        cv2.rectangle(img_bgr, (x1, pill_y1), (x1 + tw + 8, y1), color, -1)
        cv2.putText(img_bgr, label_text, (x1 + 4, y1 - 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)

        rows.append({
            "#": i + 1,
            "Finding": cls_name,
            "Confidence": f"{score:.1%}",
            "BBox": f"({x1},{y1}) → ({x2},{y2})",
        })
    return img_bgr, rows


# ── Tab 1 inference: YOLO ensemble ────────────────────────────────────────────
def run_yolo(image, conf_threshold, iou_threshold, wbf_iou, wbf_skip):
    if image is None:
        return None, pd.DataFrame(), "⚠️ Please upload an image."

    img_pil = Image.fromarray(image)
    img_w, img_h = img_pil.size
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        img_pil.save(tmp.name, quality=95)
        tmp_path = tmp.name

    all_boxes, all_scores, all_labels = [], [], []
    for cfg in YOLO_CFGS:
        preds = cfg["model"].predict(source=tmp_path, imgsz=cfg["imgsz"],
                                     conf=conf_threshold, iou=iou_threshold, verbose=False)
        res = preds[0]
        if len(res.boxes) == 0:
            all_boxes.append([]); all_scores.append([]); all_labels.append([])
            continue
        all_boxes.append(res.boxes.xyxyn.cpu().numpy().tolist())
        all_scores.append(res.boxes.conf.cpu().numpy().tolist())
        all_labels.append(res.boxes.cls.cpu().numpy().astype(int).tolist())

    fused_boxes, fused_scores, fused_labels = weighted_boxes_fusion(
        all_boxes, all_scores, all_labels,
        weights=[1, 1], iou_thr=wbf_iou, skip_box_thr=wbf_skip,
    )

    boxes_px = [[b[0]*img_w, b[1]*img_h, b[2]*img_w, b[3]*img_h] for b in fused_boxes]
    img_bgr, rows = draw_boxes(img_bgr, boxes_px, fused_labels, fused_scores,
                                list(YOLO_CLASS_NAMES.values()), YOLO_COLORS)

    annotated = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    table = pd.DataFrame(rows) if rows else pd.DataFrame(columns=["#","Finding","Confidence","BBox"])
    n = len(rows)
    if n == 0:
        summary = "✅ No findings detected above threshold."
    else:
        counts = {}
        for r in rows: counts[r["Finding"]] = counts.get(r["Finding"], 0) + 1
        summary = f"🦷 **{n} detection{'s' if n>1 else ''} found** — " + ", ".join(f"{v}× {k}" for k,v in counts.items())

    return annotated, table, summary


# ── Tab 2 inference: Faster R-CNN ────────────────────────────────────────────
def run_rcnn(image, conf_threshold):
    if image is None:
        return None, pd.DataFrame(), "⚠️ Please upload an image."
    if not RCNN_LOADED:
        return None, pd.DataFrame(), "⚠️ Faster R-CNN weights not loaded. Upload `fasterrcnn.pth` to the weights repo."

    img_pil = Image.fromarray(image).convert("RGB")
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    transform = transforms.Compose([transforms.ToTensor()])
    img_tensor = transform(img_pil).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        preds = rcnn_model(img_tensor)

    boxes  = preds[0]["boxes"].cpu().numpy()
    labels = preds[0]["labels"].cpu().numpy()
    scores = preds[0]["scores"].cpu().numpy()

    # Filter by confidence
    keep = scores >= conf_threshold
    boxes, labels, scores = boxes[keep], labels[keep], scores[keep]

    img_bgr, rows = draw_boxes(img_bgr, boxes, labels, scores,
                                RCNN_CLASS_NAMES, RCNN_COLORS)

    annotated = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    table = pd.DataFrame(rows) if rows else pd.DataFrame(columns=["#","Finding","Confidence","BBox"])
    n = len(rows)
    if n == 0:
        summary = "✅ No findings detected above threshold."
    else:
        counts = {}
        for r in rows: counts[r["Finding"]] = counts.get(r["Finding"], 0) + 1
        summary = f"🦷 **{n} detection{'s' if n>1 else ''} found** — " + ", ".join(f"{v}× {k}" for k,v in counts.items())

    return annotated, table, summary


# ── UI ─────────────────────────────────────────────────────────────────────────
with gr.Blocks(
    title="DENTO — Dental AI",
    theme=gr.themes.Base(
        primary_hue="blue",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
    ),
    css="""
    #header { text-align:center; padding:20px 0 10px; }
    #header h1 { font-size:2.2rem; font-weight:800; margin:0; }
    #header p  { color:#94a3b8; margin:6px 0 0; font-size:0.95rem; }
    """,
) as demo:

    gr.HTML("""
        <div id="header">
            <h1>🦷 DENTO</h1>
            <p>Dental X-Ray AI Detection — Two models, one interface</p>
        </div>
    """)

    with gr.Tabs():

        # ── Tab 1: AlphaDent YOLO ──────────────────────────────
        with gr.Tab("🎯 AlphaDent  (YOLOv11 Ensemble)"):
            gr.Markdown("YOLOv11 nano + small with **Weighted Box Fusion** — trained on the AlphaDent dataset.")
            with gr.Row():
                with gr.Column(scale=1):
                    yolo_input = gr.Image(label="Upload X-Ray", type="numpy", height=340)
                    with gr.Accordion("⚙️ Settings", open=False):
                        yolo_conf   = gr.Slider(0.05, 0.95, value=0.25, step=0.05, label="Confidence Threshold")
                        yolo_iou    = gr.Slider(0.10, 0.90, value=0.45, step=0.05, label="IoU Threshold (NMS)")
                        yolo_wbf    = gr.Slider(0.10, 0.90, value=0.50, step=0.05, label="WBF IoU Threshold")
                        yolo_skip   = gr.Slider(0.00, 0.30, value=0.01, step=0.01, label="WBF Skip-Box Threshold")
                    yolo_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
                with gr.Column(scale=1):
                    yolo_out_img  = gr.Image(label="Result", type="numpy", height=340)
                    yolo_summary  = gr.Markdown()
                    yolo_table    = gr.Dataframe(label="Detections", headers=["#","Finding","Confidence","BBox"], interactive=False, wrap=True)

            yolo_btn.click(fn=run_yolo,
                           inputs=[yolo_input, yolo_conf, yolo_iou, yolo_wbf, yolo_skip],
                           outputs=[yolo_out_img, yolo_table, yolo_summary])

        # ── Tab 2: Faster R-CNN ────────────────────────────────
        with gr.Tab("🏥 Dental Radiography  (Faster R-CNN)"):
            gr.Markdown("Faster R-CNN ResNet50 FPN — trained on the Dental Radiography dataset with 18 pathology classes.")
            with gr.Row():
                with gr.Column(scale=1):
                    rcnn_input = gr.Image(label="Upload X-Ray", type="numpy", height=340)
                    with gr.Accordion("⚙️ Settings", open=False):
                        rcnn_conf = gr.Slider(0.05, 0.95, value=0.50, step=0.05, label="Confidence Threshold")
                    rcnn_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
                with gr.Column(scale=1):
                    rcnn_out_img = gr.Image(label="Result", type="numpy", height=340)
                    rcnn_summary = gr.Markdown()
                    rcnn_table   = gr.Dataframe(label="Detections", headers=["#","Finding","Confidence","BBox"], interactive=False, wrap=True)

            rcnn_btn.click(fn=run_rcnn,
                           inputs=[rcnn_input, rcnn_conf],
                           outputs=[rcnn_out_img, rcnn_table, rcnn_summary])

demo.launch(server_name="0.0.0.0", server_port=7860)