Update app.py

app.py

@@ -14,6 +14,8 @@ for root, dirs, files in os.walk(base_path):
         break
 
 # Step 2: If best.pt is not found, use pre-trained weights
+model = None  # Ensure model is defined
+
 if best_path is None:
     print("Trained weights (best.pt) not found.")
     print("Using pre-trained YOLOv5 weights (yolov5s.pt) instead.")
@@ -29,6 +31,10 @@ else:
         print(f"Error loading custom model: {e}")
         model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Fallback to pre-trained model
 
+# Ensure the model was loaded properly before proceeding
+if model is None:
+    raise RuntimeError("Failed to load YOLOv5 model. Please check the weights or model path.")
+
 # Step 3: Define weapon classes to detect
 weapon_classes = ['bomb', 'gun', 'pistol', 'Automatic', 'Rifle', 'Bazooka', 
                   'Handgun', 'Knife', 'Grenade Launcher', 'Shotgun', 'SMG', 
@@ -85,3 +91,4 @@ iface = gr.Interface(
 
 # Step 5: Launch Gradio App
 iface.launch()
+v