""" 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(""" """) 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)