Update app.py

app.py

@@ -1,84 +1,68 @@
+# app.py
 import gradio as gr
-import numpy as np
 import cv2
-from ultralytics import YOLO
-from tensorflow.keras.models import load_model
-from PIL import Image, ImageDraw
+from periodontitis_detection import SimpleDentalSegmentationNoEnhance
 
 # ==========================
-# 1️⃣ Load models
+# 1️⃣ Load models once
 # ==========================
-yolo_model = YOLO("yolov8n-seg.pt")  # tooth detection
-unet_model = load_model("unet.keras", compile=False)  # CEJ/ABC segmentation
-
-# ==========================
-# 2️⃣ Helper functions
-# ==========================
-def preprocess_image(image, target_size=(256, 256)):
-    """Resize and normalize for UNet."""
-    img = np.array(image.resize(target_size)) / 255.0
-    if img.ndim == 2:
-        img = np.expand_dims(img, axis=-1)
-    return np.expand_dims(img, axis=0)
-
-def postprocess_mask(mask, original_size):
-    """Resize UNet output back to original cropped size."""
-    mask = (mask[0] > 0.5).astype(np.uint8) * 255
-    return cv2.resize(mask, original_size[::-1])
-
-def overlay_mask_on_image(image, mask):
-    """Blend segmentation mask with tooth image."""
-    color_mask = cv2.applyColorMap(mask, cv2.COLORMAP_JET)
-    blended = cv2.addWeighted(np.array(image), 0.7, color_mask, 0.3, 0)
-    return Image.fromarray(blended)
+model = SimpleDentalSegmentationNoEnhance(
+    unet_model_path="unet(10_22_25).keras",   # same filenames as your repo
+    yolo_model_path="yolov8n-seg.pt"
+)
 
 # ==========================
-# 3️⃣ Main pipeline
+# 2️⃣ Define wrapper for Gradio
 # ==========================
-def detect_and_segment(xray_image):
-    # Convert PIL → OpenCV
-    img_cv = cv2.cvtColor(np.array(xray_image), cv2.COLOR_RGB2BGR)
-
-    # Run YOLOv8 for tooth detection
-    results = yolo_model.predict(img_cv)
-    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
-
-    annotated_image = np.array(xray_image).copy()
+def detect_periodontitis(image_np):
+    """
+    Gradio sends image as a NumPy RGB array.
+    We temporarily save it to a file path since analyze_image() needs a path.
+    """
+    temp_path = "temp_input.jpg"
+    cv2.imwrite(temp_path, cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR))
 
-    segmented_teeth = []
+    # Run full pipeline
+    results = model.analyze_image(temp_path)
 
-    for i, (x1, y1, x2, y2) in enumerate(boxes):
-        # Crop individual tooth
-        crop = xray_image.crop((x1, y1, x2, y2))
+    # Convert OpenCV BGR → RGB for Gradio display
+    combined_rgb = cv2.cvtColor(results["combined"], cv2.COLOR_BGR2RGB)
+    cej_mask_rgb = cv2.cvtColor(results["cej_mask"] * 255, cv2.COLOR_GRAY2RGB)
+    abc_mask_rgb = cv2.cvtColor(results["abc_mask"] * 255, cv2.COLOR_GRAY2RGB)
 
-        # UNet segmentation
-        preprocessed = preprocess_image(crop)
-        mask = unet_model.predict(preprocessed)
-        seg_mask = postprocess_mask(mask, crop.size)
+    # Optional: summarize measurements for text output
+    summaries = []
+    for tooth in results["distance_analyses"]:
+        tooth_id = tooth["tooth_id"]
+        analysis = tooth["analysis"]
+        if analysis:
+            mean_d = analysis["mean_distance"]
+            summaries.append(f"Tooth {tooth_id}: mean={mean_d:.2f}px")
+        else:
+            summaries.append(f"Tooth {tooth_id}: no valid CEJ–ABC measurement")
 
-        # Overlay segmentation mask
-        overlay = overlay_mask_on_image(crop, seg_mask)
-        segmented_teeth.append(overlay)
+    summary_text = "\n".join(summaries) if summaries else "No detections found."
 
-        # Draw bounding box on main image
-        draw = ImageDraw.Draw(Image.fromarray(annotated_image))
-        draw.rectangle([x1, y1, x2, y2], outline="lime", width=2)
-        draw.text((x1, y1 - 10), f"Tooth {i+1}", fill="lime")
+    return combined_rgb, cej_mask_rgb, abc_mask_rgb, summary_text
 
-    return Image.fromarray(annotated_image), segmented_teeth
 
 # ==========================
-# 4️⃣ Gradio Interface
+# 3️⃣ Build Gradio Interface
 # ==========================
 demo = gr.Interface(
-    fn=detect_and_segment,
-    inputs=gr.Image(type="pil", label="Upload Dental X-Ray"),
+    fn=detect_periodontitis,
+    inputs=gr.Image(type="numpy", label="Upload Dental X-Ray"),
     outputs=[
-        gr.Image(label="Detected Teeth (YOLOv8)"),
-        gr.Gallery(label="Segmented Teeth (UNet)", show_label=True)
+        gr.Image(label="Final Annotated Image (YOLO + CEJ–ABC)"),
+        gr.Image(label="CEJ Segmentation Mask"),
+        gr.Image(label="ABC Segmentation Mask"),
+        gr.Textbox(label="Analysis Summary"),
     ],
-    title="Periodontitis Detection – YOLOv8 + UNet",
-    description="Stage 1: Tooth detection via YOLOv8. Stage 2: CEJ & ABC segmentation via UNet."
+    title="🦷 Periodontitis Detection & Analysis",
+    description=(
+        "Automatically detects teeth (YOLOv8), segments CEJ/ABC (U-Net), "
+        "and measures CEJ–ABC distances per tooth to assess bone loss."
+    ),
 )
 
 if __name__ == "__main__":