Update app.py

app.py

@@ -2,23 +2,31 @@ import gradio as gr
 from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
 from PIL import Image
+import torch
+from torch.serialization import add_safe_globals
+from ultralytics.nn.tasks import DetectionModel  # Needed for safe loading
 
-# ✅ Download your model from Hugging Face Hub
+# ✅ Allow loading DetectionModel (trusting your own model)
+add_safe_globals({'DetectionModel': DetectionModel})
+
+# ✅ Download model from Hugging Face Hub
 model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
+
+# ✅ Load model with trust
 model = YOLO(model_path)
 
 def detect_structure(image):
     results = model(image)
-    annotated_img = results[0].plot()
-    return Image.fromarray(annotated_img)
+    return Image.fromarray(results[0].plot())
 
 demo = gr.Interface(
     fn=detect_structure,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="YOLO Molecular Structure Reader",
-    description="Upload a molecular structure image. The YOLO model will detect atoms and bonds."
+    title="YOLO Molecular Structure Detector",
+    description="Upload a molecular image. YOLOv8 will detect atoms or bonds."
 )
 
 demo.launch()
 
+