Update app.py

app.py

@@ -1,13 +1,21 @@
 import gradio as gr
-from ultralytics import YOLO
 import torch
-import ultralytics.nn.tasks  # required for safe unpickling
+from ultralytics import YOLO
+
+# ✅ Allowlist required classes for safe loading
+import ultralytics.nn.tasks
+import torch.nn.modules.container
 
-# 🛠 Allow YOLOv8's custom model class to be deserialized
-torch.serialization.add_safe_globals([ultralytics.nn.tasks.DetectionModel])
+torch.serialization.add_safe_globals([
+    ultralytics.nn.tasks.DetectionModel,
+    torch.nn.modules.container.Sequential
+])
 
-# ✅ Load the trained YOLOv8 model
-model = YOLO("best.pt")  # Make sure best.pt is in the same directory
+# ✅ Load the YOLOv8 model (from trusted source like Roboflow)
+try:
+    model = YOLO("best.pt")
+except Exception as e:
+    raise RuntimeError(f"Failed to load YOLOv8 model: {e}")
 
 # 🔍 Inference function
 def detect_objects(image):
@@ -15,11 +23,11 @@ def detect_objects(image):
     annotated_image = results[0].plot()
     return annotated_image
 
-# 🎛️ Gradio interface
+# 🎛️ Gradio UI
 description = """
-📉 **Panel Fault D C Detection**  
-Upload an image to detect panel faults using a YOLOv8 model trained on Roboflow.  
-[View Dataset](https://universe.roboflow.com/naveen-kumar-9emxl/my-project-1-scnhw)
+📉 **Panel Fault D C Detection App**  
+Detect panel faults using a YOLOv8 model trained on Roboflow.  
+🔗 [Dataset Link](https://universe.roboflow.com/naveen-kumar-9emxl/my-project-1-scnhw)
 """
 
 demo = gr.Interface(