Update app.py

Show Bounding Boxes

app.py

@@ -3,41 +3,40 @@ from ultralytics import YOLO
 import numpy as np
 import cv2
 
-# Load YOLOv10 Fire + Smoke model
-model = YOLO("best.pt") 
+model = YOLO("best.pt")
 
-def detect_fire_smoke(image):
-    if image is None:
-        return "Please upload an image"
 
+def detect(image):
     img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    results = model(img)[0]
 
-    if len(results.boxes) == 0:
-        return "✔ SAFE — No Fire or Smoke Detected"
+    results = model.predict(img, conf=0.25, verbose=False)
 
-    output = []
+    annotated = results[0].plot()
 
-    for box in results.boxes:
-        cls_id = int(box.cls[0])  # YOLOv10 classes
-        conf = float(box.conf[0])
+    annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
 
-        if cls_id == 0:
-            output.append(f"🔥 FIRE DETECTED — Confidence {conf:.2f}")
-        elif cls_id == 1:
-            output.append(f"💨 SMOKE DETECTED — Confidence {conf:.2f}")
+    text = []
 
-    if not output:
-        return "✔ SAFE — No Fire or Smoke Detected"
+    for box in results[0].boxes:
+        cls = int(box.cls.item())
+        conf = float(box.conf.item())
+
+        text.append(f"{model.names[cls]} : {conf:.2%}")
+
+    if len(text) == 0:
+        text = ["✅ SAFE — No Fire or Smoke Detected"]
+
+    return annotated, "\n".join(text)
 
-    return "\n".join(output)
 
 demo = gr.Interface(
-    fn=detect_fire_smoke,
+    fn=detect,
     inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Fire & Smoke Detection (YOLOv10)",
-    description="Upload an image to detect fire or smoke."
+    outputs=[
+        gr.Image(label="Detection"),
+        gr.Textbox(label="Prediction")
+    ],
+    title="Fire & Smoke Detection"
 )
 
-demo.launch()
+demo.launch()