Update train.py

train.py

@@ -1,19 +1,14 @@
+import os
 import yaml
 from ultralytics import YOLO
-import os
-import gdown
-
-def download_pretrained_weights():
-    """Download pre-trained YOLOv8 weights"""
-    url = 'https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt'
-    gdown.download(url, 'yolov8x.pt', quiet=False)
+import streamlit as st
 
 def create_dataset_yaml():
-    """Create YAML file for dataset configuration"""
+    """Create dataset.yaml configuration."""
     data_yaml = {
-        'path': 'dataset',  # dataset root dir
-        'train': 'images/train',  # train images
-        'val': 'images/val',      # val images
+        'path': 'dataset',
+        'train': 'images/train',
+        'val': 'images/val',
         'names': {
             0: 'Commercial Airliner',
             1: 'Military Fighter',
@@ -24,30 +19,28 @@ def create_dataset_yaml():
             6: 'Unknown'
         }
     }
-    
+    os.makedirs('dataset/images/train', exist_ok=True)
+    os.makedirs('dataset/images/val', exist_ok=True)
+    os.makedirs('dataset/labels/train', exist_ok=True)
+    os.makedirs('dataset/labels/val', exist_ok=True)
+
     with open('dataset.yaml', 'w') as f:
         yaml.dump(data_yaml, f)
 
 def train_model():
-    """Train the YOLOv8 model"""
-    # Download pre-trained weights if not exists
-    if not os.path.exists('yolov8x.pt'):
-        download_pretrained_weights()
-    
-    # Create dataset configuration
-    create_dataset_yaml()
-    
-    # Initialize model
-    model = YOLO('yolov8x.pt')
-    
-    # Train the model
-    results = model.train(
-        data='dataset.yaml',
-        epochs=100,
-        imgsz=640,
-        batch=16,
-        name='aircraft_detection'
-    )
-    
-    # Export the model
-    model.export(format='onnx')
+    """Train the custom model."""
+    try:
+        create_dataset_yaml()
+        model = YOLO('yolov8n.pt')
+        model.train(
+            data='dataset.yaml',
+            epochs=100,
+            imgsz=640,
+            batch=16,
+            name='aircraft_detection'
+        )
+        model.export(format='onnx')
+        return True
+    except Exception as e:
+        st.error(f"Error during training: {str(e)}")
+        return False