Update app.py

app.py

@@ -5,23 +5,25 @@ from PIL import Image
 import torch
 import torch.serialization
 
-# ✅ Add necessary classes to safe globals
+# ✅ Add required classes for safe unpickling
 from ultralytics.nn.tasks import DetectionModel
-from torch.nn import Sequential, Conv2d
 from ultralytics.nn.modules.conv import Conv
+from torch.nn import Sequential, Conv2d, BatchNorm2d
 
-torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv, Conv2d])
+torch.serialization.add_safe_globals([
+    DetectionModel, Sequential, Conv, Conv2d, BatchNorm2d
+])
 
-# 📥 Load model
+# 📥 Load model from Hugging Face
 model_path = hf_hub_download(repo_id="Safi029/ABD-model", filename="ABD.pt")
 model = YOLO(model_path)
 
-# 🧠 Detection function
+# 🧠 Define detection function
 def detect_structure(image):
     results = model(image)
     return Image.fromarray(results[0].plot())
 
-# 🎛️ Gradio UI
+# 🎛️ Build Gradio interface
 demo = gr.Interface(
     fn=detect_structure,
     inputs=gr.Image(type="pil"),