Update app.py

app.py

@@ -2,44 +2,46 @@ import gradio as gr
 from ultralytics import YOLO
 import numpy as np
 import cv2
+from PIL import Image
 
-# Load the fire+smoke model
-model = YOLO("best.pt")
+# Load your fire+smoke model (best.pt)
+model = YOLO("best.pt")  # Make sure file exists in Space
 
 def detect_fire_smoke(image):
     if image is None:
         return "Please upload an image"
 
+    # Convert PIL → CV2
     img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
+    # Prediction
     results = model(img)[0]
 
-    # No detection
     if len(results.boxes) == 0:
-        return "✔ SAFE — NO FIRE/SMOKE DETECTED"
+        return "✔ SAFE — No Fire/Smoke Detected"
 
-    final = []
+    output = []
 
     for box in results.boxes:
-        cls_id = int(box.cls[0])
+        cls = int(box.cls[0])  # 0=fire, 1=smoke (as per model card)
         conf = float(box.conf[0])
 
-        if cls_id == 0:
-            final.append(f"🔥 FIRE DETECTED (confidence {conf:.2f})")
-        elif cls_id == 1:
-            final.append(f"💨 SMOKE DETECTED (confidence {conf:.2f})")
+        if cls == 0:
+            output.append(f"🔥 FIRE DETECTED — Confidence: {conf:.2f}")
+        elif cls == 1:
+            output.append(f"💨 SMOKE DETECTED — Confidence: {conf:.2f}")
 
-    if not final:
-        return "✔ SAFE — NO FIRE/SMOKE DETECTED"
+    if not output:
+        return "✔ SAFE — No Fire/Smoke Detected"
 
-    return "\n".join(final)
+    return "\n".join(output)
 
 demo = gr.Interface(
     fn=detect_fire_smoke,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Fire & Smoke Detection (DEMO)",
-    description="Upload an image to test detection using YOLOv10 Fire & Smoke Model."
+    title="Fire & Smoke Detection",
+    description="Upload an image to detect fire or smoke using YOLOv10 fire model."
 )
 
 demo.launch()