Update app.py

app.py

@@ -1,9 +1,12 @@
-import os
 from ultralytics import YOLO
+import os
 
-# Load a pre-trained YOLO model (adjust model type as needed)
-model = YOLO("yolo11n.yaml")  # n, s, m, l, x versions available
+# Define path to the config.yaml in your Hugging Face Space
+ROOT_DIR = '/app/data'  # Since data is uploaded under /app/data in Hugging Face Space
+config_path = os.path.join(ROOT_DIR, 'config.yaml')
 
-# Perform object detection on an image
-results = model.train(data=os.path.join(ROOT_DIR, "config.yaml"), epochs=1)  # Can also use video, directory, URL, etc.
+# 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)