Update app.py

app.py

@@ -19,8 +19,12 @@ if best_path is None:
     print("Using pre-trained YOLOv5 weights (yolov5s.pt) instead.")
     model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
 else:
-    print(f"Model weights found at: {best_path}")
-    model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
+    try:
+        print(f"Model weights found at: {best_path}")
+        model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
+    except Exception as e:
+        print(f"Error loading custom model: {e}")
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Fallback to pre-trained model
 
 # Step 3: Define weapon classes to detect
 weapon_classes = ['bomb', 'gun', 'pistol', 'Automatic', 'Rifle', 'Bazooka', 
@@ -28,7 +32,10 @@ weapon_classes = ['bomb', 'gun', 'pistol', 'Automatic', 'Rifle', 'Bazooka',
                   'Sniper', 'Sword']  # Adjust based on your dataset
 
 def detect_weapons(image):
-    results = model(image)
+    try:
+        results = model(image)
+    except Exception as e:
+        return f"Error during detection: {e}", None
 
     # Print available model class names to check for class mismatches
     model_classes = results.names  # This should give the list of class labels used by the model