Update app.py

app.py

@@ -44,44 +44,49 @@ def base64_to_pil(b64_str):
 
 
 # --- Main detection ---
-def detect_gingivitis(image, conf=0.25, iou=0.5):
-    image = preprocess(image)
-
-    sw_res = model_swelling.predict(image, conf=conf, iou=iou)
-    rd_res = model_redness.predict(image, conf=conf, iou=iou)
-    bl_res = model_bleeding.predict(image, conf=conf, iou=iou)
-
-    # Draw boxes only (no labels)
-    img_sw = sw_res[0].plot(labels=False)[:, :, ::-1]
-    img_rd = rd_res[0].plot(labels=False)[:, :, ::-1]
-    img_bl = bl_res[0].plot(labels=False)[:, :, ::-1]
-
-    # Encode images to Base64 (for backend API consumption)
-    sw_b64 = np_to_base64(img_sw)
-    rd_b64 = np_to_base64(img_rd)
-    bl_b64 = np_to_base64(img_bl)
-
-    # Convert Base64 back to PIL for Gradio display
-    sw_pil = base64_to_pil(sw_b64)
-    rd_pil = base64_to_pil(rd_b64)
-    bl_pil = base64_to_pil(bl_b64)
-
-    # Determine diagnosis
-    has_sw = len(sw_res[0].boxes) > 0
-    has_rd = len(rd_res[0].boxes) > 0
-    has_bl = len(bl_res[0].boxes) > 0
-
-    if has_sw and has_rd and has_bl:
-        diagnosis = (
-            "🦷 You have gingivitis.\n\n"
-            "For accurate assessment and guidance, we recommend visiting your dentist.\n"
-            "If you have a periapical X-ray, you may try the *Detect Periodontitis* tool."
-        )
-    else:
-        diagnosis = "🟢 You don't have gingivitis."
-
-    # Return PIL for Gradio + Base64 is available for backend
-    return [sw_pil, rd_pil, bl_pil, diagnosis]
+def detect_gingivitis_safe(image, conf=0.25, iou=0.5):
+    try:
+        if image is None:
+            return [None, None, None, "❌ No image uploaded"]
+
+        # Lazy-load models
+        model_swelling = get_swelling_model()
+        model_redness = get_redness_model()
+        model_bleeding = get_bleeding_model()
+
+        # Preprocess image
+        image = preprocess(image)
+
+        # Run models
+        sw_res = model_swelling.predict(image, conf=conf, iou=iou)
+        rd_res = model_redness.predict(image, conf=conf, iou=iou)
+        bl_res = model_bleeding.predict(image, conf=conf, iou=iou)
+
+        # Convert to images
+        img_sw = sw_res[0].plot(labels=False)[:, :, ::-1]
+        img_rd = rd_res[0].plot(labels=False)[:, :, ::-1]
+        img_bl = bl_res[0].plot(labels=False)[:, :, ::-1]
+
+        sw_pil = base64_to_pil(np_to_base64(img_sw))
+        rd_pil = base64_to_pil(np_to_base64(img_rd))
+        bl_pil = base64_to_pil(np_to_base64(img_bl))
+
+        # Determine diagnosis
+        has_sw = len(sw_res[0].boxes) > 0
+        has_rd = len(rd_res[0].boxes) > 0
+        has_bl = len(bl_res[0].boxes) > 0
+
+        if has_sw and has_rd and has_bl:
+            diagnosis = "🦷 You have gingivitis."
+        else:
+            diagnosis = "🟢 You don't have gingivitis."
+
+        return [sw_pil, rd_pil, bl_pil, diagnosis]
+
+    except Exception as e:
+        # Catch all errors and return a friendly message
+        return [None, None, None, f"❌ Error during processing: {str(e)}"]
+
 
 
 # --- Gradio Interface ---