Update app.py

app.py

@@ -4,12 +4,21 @@ import os
 # Define the correct path to config.yaml (in the root directory)
 config_path = './config.yaml'  # Adjust based on the actual path to your config.yaml
 
-
 # Load YOLO model
 model = YOLO("yolo11n.yaml")  # You can choose a different model type like yolo5n, yolo6n, etc.
 
 # Train the model
 results = model.train(data=config_path, epochs=1)
-print("Model saved to:", 'runs/detect/train/weights/best.pt')
-model.save('path_to_save/best.pt')
 
+# Define the save directory
+save_dir = './runs/detect/train/weights'
+
+# Create directory if it doesn't exist
+if not os.path.exists(save_dir):
+    os.makedirs(save_dir)
+
+# Save the model
+model.save(os.path.join(save_dir, 'best.pt'))
+
+# Print confirmation
+print("Model saved to:", os.path.join(save_dir, 'best.pt'))