ToothFairy/train_valid_test_split.py

from dagshub.streaming import install_hooks
install_hooks()

import dagshub
from dagshub.data_engine import datasources

import logging

from utils.data import create_splits

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

# Environment Variables
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"

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)

    logger.info('Create train/valid/test splits')
    md = ds.all().dataframe
    md = create_splits(md)

    # Upload the metadata
    logger.info('Uploading the metadata to Data Engine')
    dagshub.common.config.dataengine_metadata_upload_batch_size = 50
    ds.upload_metadata_from_dataframe(md, path_column="path")

if __name__ == '__main__':
    main()