Update app.py

app.py

@@ -3,23 +3,24 @@ from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import torch
+import torch.serialization
 
-# Download the model file from Hugging Face
-model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
+# 🛡️ ALLOW custom classes in checkpoint (required for torch >=2.6)
+from ultralytics.nn.tasks import DetectionModel
+torch.serialization.add_safe_globals([DetectionModel])
 
-# Explicitly load with weights_only=False
-ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
+# 📥 Download model from Hugging Face repo
+model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
 
-# Now manually initialize YOLO model from checkpoint
-model = YOLO()
-model.model.load_state_dict(ckpt['model'].float().state_dict())
+# ✅ Load full model (not weights only)
+model = YOLO(model_path)
 
-# Function to run detection
+# 🔍 Detection function
 def detect_structure(image):
     results = model(image)
     return Image.fromarray(results[0].plot())
 
-# Gradio interface
+# 🎛️ Gradio UI
 demo = gr.Interface(
     fn=detect_structure,
     inputs=gr.Image(type="pil"),