ToothFairy/enrich_metadata.py

from dagshub.streaming import install_hooks
install_hooks()

import dagshub
from dagshub.data_engine import datasources

import logging

from utils.data import DataFunctions

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)

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

    logger.info('Converting the datasource to a pandas Dataframe')
    md = ds.all().dataframe

    # Enrich metadata
    logger.info('Enriching the metadata')
    enriched_md = md.apply(lambda x: dataset_func.create_metadata(x), axis=1)

    # 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(enriched_md, path_column="path")

if __name__ == '__main__':
    main()