Update app.py

app.py

@@ -49,7 +49,12 @@ print(f"CUDA available: {torch.cuda.is_available()}")
 # Load custom YOLO model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
-model = YOLO('/mnt/data/best.pt').to(device)  # Updated model path
+
+# Model paths
+model_path = '/mnt/data/best.pt'  # Path to the custom-trained model
+
+# Load the model
+model = YOLO(model_path).to(device)  # Updated model path
 if device == "cuda":
     model.half()  # Use half-precision (FP16)
 print(f"Model classes: {model.names}")