Update app.py

app.py

@@ -1,20 +1,28 @@
 import gradio as gr
+import torch
+from torch.serialization import add_safe_globals
 from ultralytics import YOLO
 import numpy as np
 import cv2
 from PIL import Image
+import ultralytics.nn.tasks as yolov8_tasks
 
-# Load your fire+smoke model (best.pt)
-model = YOLO("best.pt")  # Make sure file exists in Space
+# ---------------------------
+# Allow YOLO model class for PyTorch 2.6+
+# ---------------------------
+add_safe_globals([yolov8_tasks.DetectionModel])
+
+# ---------------------------
+# Load model safely
+# ---------------------------
+model = YOLO("best.pt")   # Works after safe global patch
 
 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]
 
     if len(results.boxes) == 0:
@@ -23,7 +31,7 @@ def detect_fire_smoke(image):
     output = []
 
     for box in results.boxes:
-        cls = int(box.cls[0])  # 0=fire, 1=smoke (as per model card)
+        cls = int(box.cls[0])
         conf = float(box.conf[0])
 
         if cls == 0:
@@ -41,7 +49,7 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="Fire & Smoke Detection",
-    description="Upload an image to detect fire or smoke using YOLOv10 fire model."
+    description="Upload an image to detect fire or smoke using YOLOv10 model."
 )
 
 demo.launch()