ToothFairy/train_yolo.py

from dagshub.streaming import install_hooks
install_hooks()

import dagshub
from dagshub.data_engine import datasources

import logging
import mlflow
import ultralytics

from utils.data import DataFunctions
from utils.dagshub_yolo_cb import custom_callbacks_fn

logger = logging.getLogger('root')
logger.setLevel(logging.INFO)

# Environment Variables
DAGSHUB_USER = "yonomitt"
DAGSHUB_REPO_OWNER = "yonomitt"
DAGSHUB_REPO="ToothFairy"
DAGSHUB_FULL_REPO=DAGSHUB_REPO_OWNER + "/" + DAGSHUB_REPO
DATASOURCE_NAME = "Tooth-Segmentation"
DATASOURCE_PATH = "s3://tooth-dataset/data"
ANNOTATION_FILE = "s3://tooth-dataset/tooth_segmentation.json"
MLFLOW_PROJECT = "Default"

def get_or_create_datasource(name):
    try:
        ds = datasources.get_datasource(repo=DAGSHUB_FULL_REPO, name=name)
    except:
        ds = datasources.create(repo=DAGSHUB_FULL_REPO, name=name, path=DATASOURCE_PATH)
    return ds

def main():
    logger.info('Getting or creating the datasource')
    ds = get_or_create_datasource(DATASOURCE_NAME)

    dataset_func = DataFunctions(annotation_file=ANNOTATION_FILE, yolo_dir='generated/yolo_data', label_type='segmentation')

    logger.info('Converting the datasource to a YOLO-compatible dataset')

    dataset_func.create_yolo_v8_dataset_yaml(ds, download=False)

    logger.info('Monkey-patching ultralytics')
    ultralytics.utils.callbacks.add_integration_callbacks = custom_callbacks_fn

    dagshub.init(repo_name=DAGSHUB_REPO, repo_owner=DAGSHUB_USER)

    # Load a model
    model = ultralytics.YOLO('yolov8l-seg.pt', task='segment')  # load a pretrained model (recommended for training)

    with mlflow.start_run():

        # Train the model
        model.train(data='custom_yolo.yaml', epochs=100, imgsz=640, device=0, project=MLFLOW_PROJECT)

if __name__ == '__main__':
    main()