Kaggleをする際の環境を区別してPATHを設定するコード

import os
import sys

def classify_env(TRIAL_NAME: str, NOW: str):

    # Colab
    if 'COLAB_GPU' in set(os.environ.keys()):
        DATA_DIR = '/content/drive/MyDrive/Kaggle/Cassava/data'
        OUTPUT_DIR = f'/content/drive/MyDrive/Kaggle/Cassava/output/{NOW}-{TRIAL_NAME}'
        os.makedirs(OUTPUT_DIR, exist_ok=True)

        sys.path.append(f'{DATA_DIR}/pytorch-image-models-master')

        print('Use COLAB')

    # kaggle
    elif 'KAGGLE_URL_BASE' in set(os.environ.keys()):
        DATA_DIR = '../input/cassava-leaf-disease-classification/'
        OUTPUT_DIR = './'

        sys.path.append('../input/pytorch-image-models/pytorch-image-models-master')

        print('Use Kaglle')


    # GCP
    elif 'GSETTINGS_SCHEMA_DIR' in set(os.environ.keys()):
        DATA_DIR = '/home/jupyter/data'
        OUTPUT_DIR = f'/home/jupyter/output/{NOW}-{TRIAL_NAME}'
        os.makedirs(OUTPUT_DIR, exist_ok=True)

        sys.path.append('../data/pytorch-image-models-master')

        print('Use GCP')
    
    print(f'DATA_DIR: {DATA_DIR}')
    print(f'OUTPUT_DIR: {OUTPUT_DIR}')
    print(f'NOW: {NOW}')

    return DATA_DIR, OUTPUT_DIR
    
DATA_DIR, OUTPUT_DIR = classify_env(TRIAL_NAME, NOW)