Annotation-Loop-Showcase/scripts/model_training.py

"""
Training a model on the data from a Data Engine datasource

NOTE: pulling the repository with dvc is not required for this,
since the images + annotations are all received from Data Engine

Furthermore, if you did any annotation work in LabelStudio, the annotations might have diverged
"""
import logging
import os
from pathlib import Path

import dagshub
import dagshub.mlflow
import mlflow
from dagshub.common.determine_repo import determine_repo
from dagshub.data_engine import datasources
from dagshub.data_engine.model.datasource import Datasource
from dagshub_annotation_converter.image.exporters import YoloExporter
from dagshub_annotation_converter.image.importers import DagshubDatasourceImporter
from ultralytics import YOLO

logger = logging.getLogger(__name__)

DATASOURCE_NAME = "COCO128"  # Name of the Data Engine datasource


def build_training_data(ds: Datasource, training_dir: str) -> Path:
    """
    Download the datapoints + generate the annotations from a Data Engine datasource, returns a path to the directory
    """
    target_dir = Path(training_dir)
    target_dir.mkdir(parents=True, exist_ok=True)
    # Download the datapoints
    query_result = ds.all()
    query_result.download_files(target_dir=target_dir)

    # Reimport the annotations
    importer = DagshubDatasourceImporter(query_result)
    exporter = YoloExporter(data_dir=target_dir, meta_file=target_dir / "yolo.yaml",
                            annotation_type="bbox")

    proj = importer.parse()
    exporter.export(proj)

    return target_dir


def get_datasource(ds_name: str) -> Datasource:
    current_repo, branch = determine_repo()
    ds = datasources.get_datasource(current_repo.full_name, ds_name)
    return ds

def train_model():

    model = YOLO("yolov8n.pt")
    model.train(data="yolo.yaml", device="mps", epochs=3)
    # After this the model should be uploaded as an artifact to the MLflow run


def main():
    current_repo, branch = determine_repo()
    # Run dagshub.init to log training metrics into MLflow
    dagshub.init(repo_owner=current_repo.owner, repo_name=current_repo.repo_name, mlflow=True)
    # Turn on autolog so YOLO training gets logged
    mlflow.autolog()
    # Turn on reliability patch so mlflow logging failing doesn't crash training
    dagshub.mlflow.patch_mlflow()
    with mlflow.start_run():
        ds = get_datasource(DATASOURCE_NAME)
        training_dir = build_training_data(ds, "training/")
        os.chdir(training_dir)
        train_model()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # Disable noisy GQL logs
    from gql.transport.requests import log as requests_logger

    requests_logger.setLevel(logging.WARNING)

    # Change current directory to the root of the repository
    os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
    main()