Update app.py

app.py

@@ -3,22 +3,23 @@ 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
 
-# ✅ Allow loading DetectionModel (trusting your own model)
-add_safe_globals({'DetectionModel': DetectionModel})
-
-# ✅ Download model from Hugging Face Hub
+# Download the model file from Hugging Face
 model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
 
-# ✅ Load model with trust
-model = YOLO(model_path)
+# Explicitly load with weights_only=False
+ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+
+# Now manually initialize YOLO model from checkpoint
+model = YOLO()
+model.model.load_state_dict(ckpt['model'].float().state_dict())
 
+# Function to run detection
 def detect_structure(image):
     results = model(image)
     return Image.fromarray(results[0].plot())
 
+# Gradio interface
 demo = gr.Interface(
     fn=detect_structure,
     inputs=gr.Image(type="pil"),