Update app.py

app.py

@@ -1,48 +1,20 @@
 import gradio as gr
-import torch
 from ultralytics import YOLO
 
-# ✅ Import modules needed for safe unpickling
-import ultralytics.nn.tasks
-import ultralytics.nn.modules.conv
-import torch.nn.modules.container
-import torch.nn.modules.conv
+# ✅ Load model directly from local file (same folder as app.py)
+model = YOLO("best.pt")
 
-# ✅ Add all required classes to the safe globals list
-torch.serialization.add_safe_globals([
-    ultralytics.nn.tasks.DetectionModel,
-    ultralytics.nn.modules.conv.Conv,
-    torch.nn.modules.container.Sequential,
-    torch.nn.modules.conv.Conv2d,
-])
-
-# ✅ Load YOLOv8 model
-try:
-    print("⚙️ Loading YOLOv8 model...")
-    model = YOLO("best.pt")
-except Exception as e:
-    raise RuntimeError(f"❌ Failed to load YOLOv8 model:\n{e}")
-
-# 🧠 Inference function
-def detect_objects(image):
+def predict(image):
     results = model(image)
-    annotated_image = results[0].plot()
-    return annotated_image
+    return results[0].plot()
 
-# 🎛️ Gradio interface
-description = """
-📉 **Panel Fault D C Detection**  
-Upload an image to detect panel faults using your custom-trained YOLOv8 model.  
-🔗 [View Dataset](https://universe.roboflow.com/naveen-kumar-9emxl/my-project-1-scnhw)
-"""
-
-demo = gr.Interface(
-    fn=detect_objects,
+iface = gr.Interface(
+    fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Panel Fault D C",
-    description=description
+    outputs=gr.Image(type="pil", label="Detected Objects"),
+    title="Roboflow YOLOv8 Detector",
+    description="Upload an image to detect objects using your YOLOv8 model trained with Roboflow."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+iface.launch()
+