Update app.py

app.py

@@ -36,43 +36,30 @@ def preprocess(image):
 def detect_gingivitis(image, conf=0.4, iou=0.5):
     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)
 
-    # >> Create annotated PIL images <<
-    img_sw = Image.fromarray(sw_res[0].plot()[..., ::-1])
-    img_rd = Image.fromarray(rd_res[0].plot()[..., ::-1])
-    img_bl = Image.fromarray(bl_res[0].plot()[..., ::-1])
+    # Convert YOLO output → numpy RGB
+    img_sw = sw_res[0].plot()[:, :, ::-1]   # BGR → RGB
+    img_rd = rd_res[0].plot()[:, :, ::-1]
+    img_bl = bl_res[0].plot()[:, :, ::-1]
 
-    # >> Save to temporary files <<
-    sw_path = save_temp_file(img_sw)
-    rd_path = save_temp_file(img_rd)
-    bl_path = save_temp_file(img_bl)
-
-    # >> 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
 
-    # ------------------------------
-    #  UPDATED: CUSTOM DIAGNOSIS
-    # ------------------------------
     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 for a professional evaluation "
-            "or second opinion. Gingivitis can progress into periodontitis if left untreated.\n\n"
-            "If you have a periapical X-ray available, you can use the **Detect Periodontitis** feature for further analysis.\n\n"
-            "Would you like to proceed?"
+            "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.**"
+        diagnosis = "🟢 You don't have gingivitis."
+
+    return [img_sw, img_rd, img_bl, diagnosis]
 
-    # *** CRITICAL ***
-    # Must return FILE PATHS + diagnosis string in ONE LIST
-    return [sw_path, rd_path, bl_path, diagnosis]
 
 # Gradio Interface
 interface = gr.Interface(
@@ -83,9 +70,9 @@ interface = gr.Interface(
         gr.Slider(0, 1, value=0.5, step=0.05, label="NMS IoU Threshold"),
     ],
     outputs=[
-        gr.File(label="Swelling Detection"),
-        gr.File(label="Redness Detection"),
-        gr.File(label="Bleeding Detection"),
+        gr.Image(label="Swelling Detection"),
+        gr.Image(label="Redness Detection"),
+        gr.Image(label="Bleeding Detection"),
         gr.Textbox(label="Diagnosis")
     ],
     title="Gingivitis Detection"