Create app.py

app.py

@@ -0,0 +1,85 @@
+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
+
+# ==========================
+# 1️⃣ Load models
+# ==========================
+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)
+
+# ==========================
+# 3️⃣ Main pipeline
+# ==========================
+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()
+
+    segmented_teeth = []
+
+    for i, (x1, y1, x2, y2) in enumerate(boxes):
+        # Crop individual tooth
+        crop = xray_image.crop((x1, y1, x2, y2))
+
+        # UNet segmentation
+        preprocessed = preprocess_image(crop)
+        mask = unet_model.predict(preprocessed)
+        seg_mask = postprocess_mask(mask, crop.size)
+
+        # Overlay segmentation mask
+        overlay = overlay_mask_on_image(crop, seg_mask)
+        segmented_teeth.append(overlay)
+
+        # 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 Image.fromarray(annotated_image), segmented_teeth
+
+# ==========================
+# 4️⃣ Gradio Interface
+# ==========================
+demo = gr.Interface(
+    fn=detect_and_segment,
+    inputs=gr.Image(type="pil", label="Upload Dental X-Ray"),
+    outputs=[
+        gr.Image(label="Detected Teeth (YOLOv8)"),
+        gr.Gallery(label="Segmented Teeth (UNet)", show_label=True)
+    ],
+    title="Periodontitis Detection – YOLOv8 + UNet",
+    description="Stage 1: Tooth detection via YOLOv8. Stage 2: CEJ & ABC segmentation via UNet."
+)
+
+if __name__ == "__main__":
+    demo.launch()