Update app.py

app.py

@@ -2,18 +2,13 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image, ImageOps, ImageEnhance
 import numpy as np
-import tempfile
+import tempfile, io, base64
 
 # Load models
 model_swelling = YOLO("models/swelling/swelling_final.pt")
 model_redness = YOLO("models/redness/redness_final.pt")
 model_bleeding = YOLO("models/bleeding/bleeding_final.pt")
 
-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):
     if isinstance(image, np.ndarray):
@@ -33,6 +28,15 @@ def preprocess(image):
 
     return image
 
+
+def np_to_base64(img_np, format="JPEG"):
+    """Convert a numpy RGB image to Base64 string."""
+    pil_img = Image.fromarray(img_np)
+    buffer = io.BytesIO()
+    pil_img.save(buffer, format=format)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
 def detect_gingivitis(image, conf=0.4, iou=0.5):
     image = preprocess(image)
 
@@ -41,10 +45,15 @@ def detect_gingivitis(image, conf=0.4, iou=0.5):
     bl_res = model_bleeding.predict(image, conf=conf, iou=iou)
 
     # Convert YOLO output → numpy RGB
-    img_sw = sw_res[0].plot()[:, :, ::-1]   # BGR → RGB
+    img_sw = sw_res[0].plot()[:, :, ::-1]  # BGR → RGB
     img_rd = rd_res[0].plot()[:, :, ::-1]
     img_bl = bl_res[0].plot()[:, :, ::-1]
 
+    # Encode images to Base64
+    sw_b64 = np_to_base64(img_sw)
+    rd_b64 = np_to_base64(img_rd)
+    bl_b64 = np_to_base64(img_bl)
+
     has_sw = len(sw_res[0].boxes) > 0
     has_rd = len(rd_res[0].boxes) > 0
     has_bl = len(bl_res[0].boxes) > 0
@@ -58,7 +67,8 @@ def detect_gingivitis(image, conf=0.4, iou=0.5):
     else:
         diagnosis = "🟢 You don't have gingivitis."
 
-    return [img_sw, img_rd, img_bl, diagnosis]
+    # Return Base64 images + diagnosis text
+    return [sw_b64, rd_b64, bl_b64, diagnosis]
 
 
 # Gradio Interface
@@ -70,12 +80,12 @@ interface = gr.Interface(
         gr.Slider(0, 1, value=0.5, step=0.05, label="NMS IoU Threshold"),
     ],
     outputs=[
-        gr.Image(label="Swelling Detection"),
-        gr.Image(label="Redness Detection"),
-        gr.Image(label="Bleeding Detection"),
+        gr.Image(label="Swelling Detection", type="auto"),
+        gr.Image(label="Redness Detection", type="auto"),
+        gr.Image(label="Bleeding Detection", type="auto"),
         gr.Textbox(label="Diagnosis")
     ],
     title="Gingivitis Detection"
 )
 
-interface.launch()
+interface.launch()