train_segmentation.py

from ultralytics import YOLO

def train_rice_segmentation():
    # Load the base YOLOv8 Nano Segmentation model
    model = YOLO('yolov8n-seg.pt')

    # Train the model
    # Using the existing data.yaml which contains the polygon annotations
    results = model.train(
        data='dataset/data.yaml',
        epochs=50,
        imgsz=640,   
        batch=8,     
        name='rice_quality_seg_v1',
        project='runs/segment',
        device=0,    # Use GPU
        workers=0    # Avoid multiprocessing overhead
    )

    # Export the best model to ONNX for production deployment
    path = model.export(format='onnx')
    print(f"Segmentation model exported to: {path}")

if __name__ == "__main__":
    train_rice_segmentation()