Update app.py

app.py

@@ -17,7 +17,10 @@ for root, dirs, files in os.walk(base_path):
 if best_path is None:
     print("Trained weights (best.pt) not found.")
     print("Using pre-trained YOLOv5 weights (yolov5s.pt) instead.")
-    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
+    try:
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
+    except Exception as e:
+        print(f"Error loading pre-trained YOLOv5 model: {e}")
 else:
     try:
         print(f"Model weights found at: {best_path}")
@@ -37,8 +40,8 @@ def detect_weapons(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
+    # Check available model class names
+    model_classes = results.names
     print("Model class names:", model_classes)
 
     # Filter detections by confidence threshold (0.5 or higher)
@@ -47,8 +50,6 @@ def detect_weapons(image):
 
     # Get the detected classes with high confidence
     detected_classes = filtered_results['name'].unique()
-
-    # Print detected classes for debugging
     print("Detected classes:", detected_classes)
 
     # Check if any of the detected objects are weapons