Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from ultralytics import YOLO
-from PIL import Image
-from PIL import ImageOps, ImageEnhance
+from PIL import Image, ImageOps, ImageEnhance
 import numpy as np
 import tempfile
 
@@ -10,61 +9,58 @@ model_swelling = YOLO("models/swelling/best.pt")
 model_redness = YOLO("models/redness/best.pt")
 model_bleeding = YOLO("models/bleeding/best.pt")
 
-def save_temp_image(img):
-    """Save PIL image to temp file and return the file path."""
+def save_temp_file(img):
+    """Save PIL image to a temporary file and return the file path."""
     temp = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
     img.save(temp.name, format="JPEG")
     return temp.name
 
-def preprocess_image(image, max_size=1024, contrast_factor=1.05):
-    if image is None:
-        return image
-
+def preprocess(image):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
 
-    image = ImageOps.exif_transpose(image)
-    image = image.convert("RGB")
+    image = ImageOps.exif_transpose(image).convert("RGB")
 
+    # Resize if needed
     w, h = image.size
-    max_side = max(w, h)
-    if max_side > max_size:
-        scale = max_size / float(max_side)
-        new_size = (int(w * scale), int(h * scale))
-        image = image.resize(new_size, Image.LANCZOS)
+    max_dim = max(w, h)
+    if max_dim > 1024:
+        scale = 1024 / max_dim
+        image = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
 
-    enhancer = ImageEnhance.Contrast(image)
-    image = enhancer.enhance(contrast_factor)
+    # Light contrast boost
+    image = ImageEnhance.Contrast(image).enhance(1.05)
 
     return image
 
-def detect_gingivitis(image, conf_threshold=0.5, iou_threshold=0.5):
-    image = preprocess_image(image)
+def detect_gingivitis(image, conf=0.4, iou=0.5):
+    image = preprocess(image)
 
-    # Swelling
-    res_swelling = model_swelling.predict(image, conf=conf_threshold, iou=iou_threshold)
-    swelling_img = Image.fromarray(res_swelling[0].plot()[..., ::-1])
-    swelling_path = save_temp_image(swelling_img)
+    # >> 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)
 
-    # Redness
-    res_redness = model_redness.predict(image, conf=conf_threshold, iou=iou_threshold)
-    redness_img = Image.fromarray(res_redness[0].plot()[..., ::-1])
-    redness_path = save_temp_image(redness_img)
+    # >> 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])
 
-    # Bleeding
-    res_bleeding = model_bleeding.predict(image, conf=conf_threshold, iou=iou_threshold)
-    bleeding_img = Image.fromarray(res_bleeding[0].plot()[..., ::-1])
-    bleeding_path = save_temp_image(bleeding_img)
+    # >> 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 text
-    has_swelling = len(res_swelling[0].boxes) > 0
-    has_redness = len(res_redness[0].boxes) > 0
-    has_bleeding = len(res_bleeding[0].boxes) > 0
+    # >> 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
 
-    diagnosis = "Gingivitis Detected" if (has_swelling and has_redness and has_bleeding) else "No Gingivitis Detected"
+    diagnosis = "Gingivitis Detected" if (has_sw and has_rd and has_bl) else "No Gingivitis Detected"
 
-    # Must return EXACTLY what Flask expects:
-    return [swelling_path, redness_path, bleeding_path, diagnosis]
+    # *** CRITICAL ***
+    # Must return FILE PATHS + diagnosis string in ONE LIST
+    return [sw_path, rd_path, bl_path, diagnosis]
 
 # Gradio Interface
 interface = gr.Interface(
@@ -78,10 +74,9 @@ interface = gr.Interface(
         gr.File(label="Swelling Detection"),
         gr.File(label="Redness Detection"),
         gr.File(label="Bleeding Detection"),
-        gr.Textbox(label="Gingivitis Diagnosis")
+        gr.Textbox(label="Diagnosis")
     ],
-    title="Gingivitis Detection (3 Models)",
-    description="Detect swelling, redness, and bleeding using 3 YOLO models.",
+    title="Gingivitis Detection"
 )
 
 interface.launch()