import os import subprocess import sys from tqdm import tqdm PD4Dval = [ "000004","000054","000075","000078","000101","000104","000109","000144","000150","000172","000178", "000190","000263","000303","000309","000316","000610","000625","000676","000683","000686","000700", "000716","000739","000936","000938","000950","001015","001059","001061","001064","001082","001083", "001335","001344","001346","001357","001389","001411","001416","001435","001669","001709","001715", "001719","001728","001773","001775","001787","001825","001887","001922","001937","002023","002035", "002070","002073","002088","002096","002138","002142"] PD4Dtest = [ "000048","000072","000074","000106","000165","000167","000173","000181","000294","000298","000310", "000315","000337","000362","000365","000444","000452","000581","000626","000651","000663","000687", "000730","000743","000839","000881","000901","000974","001010","001113","001126","001137","001235", "001347","001349","001360","001365","001369","001370","001420","001470","001483","001505","001508", "001609","001617","001664","001698","001721","001755","001783","001791","001803","001809","001885", "001946","001964","002010","002104","002120","002131"] def list_s3(s3_path): proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True) (data, _) = proc.communicate() data = data.decode().replace('/\n', ' ').replace('\n', ' ').replace(' ', ' ').split(' ') data = [o for o in data if len(o) > 0 and o not in ['PRE']] proc.kill() return data def list_s3_recursive(s3_path): proc = subprocess.Popen(f'aws s3 ls {s3_path}/ --recursive', stdout=subprocess.PIPE, shell=True) (data, _) = proc.communicate() data = data.decode().replace('/\n', ' ').replace('\n', ' ').replace(' ', ' ').split(' ') data = [o for o in data if len(o) > 0 and o not in ['PRE']] proc.kill() return data s3_root = 's3://tri-ml-sandbox-16011-us-west-2-datasets' s3_folder = 'cv_datasets/webdatasets' def create_indices_single(path, nested): if nested > 0: folders = list_s3(path) for f in folders: create_indices_single(f'{path}/{f}', nested - 1) return path = path.replace(f'{s3_root}/{s3_folder}/', '') if path.endswith('/'): path = path[:-1] name = f"{path.replace('/', '___')}.txt" path = f'{s3_root}/{s3_folder}/{path}' print() print('###########################################') print('######', path) print('###########################################') print() folders = list_s3_recursive(f'{path}') folders = [f for f in folders if f.endswith('tar')] with open(name, 'w') as file: for f in tqdm(folders): file.write(f'{s3_root}/{f}\n') os.system(f'aws s3 cp {name} {path.replace("tarfiles","indices")}/split_all.txt --quiet') os.system(f'rm {name}') ######################################################################################### def save_split(split, func): name_split = name.replace('.txt', f'_{split}.txt') total = 0 with open(name_split, 'w') as file_split: for f in tqdm(folders, ncols=128): seq = f.split('/')[-2] if func(seq): file_split.write(f'{s3_root}/{f}\n') total += 1 print(f'Saving split {split} : {total}') os.system(f'aws s3 cp {name_split} {path.replace("tarfiles","indices")}/split_{split}.txt --quiet') os.system(f'rm {name_split}') ######################################################################################### dataset = folders[0].split('/')[3] ####################################### ArgoVerse2sync def ArgoVerse2sync_train(sequence): return sequence[0] != '1' def ArgoVerse2sync_val(sequence): return sequence[0] == '1' if dataset == 'ArgoVerse2sync': save_split('train', ArgoVerse2sync_train) save_split('val', ArgoVerse2sync_val) ####################################### AssemblyHands def AssemblyHands_train(sequence): return int(sequence.split('_')[-1]) < 18 def AssemblyHands_val(sequence): return int(sequence.split('_')[-1]) >= 18 if dataset == 'AssemblyHands': save_split('train', AssemblyHands_train) save_split('val', AssemblyHands_val) ####################################### DDAD # Sequences numbered >= 150 are val, rest is train (official split) def DDAD_train(sequence): return int(sequence) < 150 def DDAD_val(sequence): return int(sequence) >= 150 if dataset == 'DDAD': save_split('train', DDAD_train) save_split('val', DDAD_val) ####################################### DL3DV # Sequences with the leftmost digit == 2 are val, rest is train def DL3DV_train(sequence): return sequence.split('__')[1][0] != '2' def DL3DV_val(sequence): return sequence.split('__')[1][0] == '2' if dataset == 'DL3DV': save_split('train', DL3DV_train) save_split('val', DL3DV_val) ####################################### DROIDCalib # Sequences with the rightmost digit == 9 are val, rest is train def DROIDCalib_train(sequence): return sequence[-1] != '9' def DROIDCalib_val(sequence): return sequence[-1] == '9' if dataset == 'DROIDCalib': save_split('train', DROIDCalib_train) save_split('val', DROIDCalib_val) ####################################### DROIDCalibV2 # Sequences with the rightmost digit == 9 are val, rest is train def DROIDCalibV2_train(sequence): return sequence[-1] != '9' def DROIDCalibV2_val(sequence): return sequence[-1] == '9' if dataset == 'DroidCalibV2': save_split('train', DROIDCalibV2_train) save_split('val', DROIDCalibV2_val) ####################################### EgoExo4D # # Proposed split separation # def EgoExo4D_train(sequence): # return not(EgoExo4D_val(sequence)) # def EgoExo4D_val(sequence): # heldout_institution = 'fair_' in sequence.lower() or \ # 'nus_' in sequence.lower() # heldout_activity = '_cpr_' in sequence.lower() or \ # '_guitar_' in sequence.lower() # heldout_pair = 'iiith_piano_' in sequence.lower() or \ # 'uniandes_basketball_' in sequence.lower() or \ # 'utokyo_soccer_' in sequence.lower() # heldout_any = heldout_institution or heldout_activity or heldout_pair # if heldout_any: # return True # else: # # last digit BEFORE the underscore to denote CAPTURE (because VERY last digit is TAKE) # return sequence.split('_')[-2][-1] == '2' # if dataset == 'EgoExo4D': # save_split('train3', EgoExo4D_train) # save_split('val3', EgoExo4D_val) ####################################### EgoExo4D # Proposed split separation (see google excel sheet) def EgoExo4D_train(sequence): return not(EgoExo4D_val_ood(sequence)) and \ not(EgoExo4D_val_id_capture(sequence)) def EgoExo4D_val_ood_site(sequence): heldout_site = 'fair_' in sequence.lower() or \ 'nus_' in sequence.lower() return heldout_site def EgoExo4D_val_ood_activity(sequence): heldout_activity = '_cpr_' in sequence.lower() or \ '_guitar_' in sequence.lower() return heldout_activity def EgoExo4D_val_ood_pair(sequence): heldout_pair = 'iiith_piano_' in sequence.lower() or \ 'uniandes_basketball_' in sequence.lower() or \ 'utokyo_soccer_' in sequence.lower() return heldout_pair def EgoExo4D_val_ood(sequence): return EgoExo4D_val_ood_site(sequence) or \ EgoExo4D_val_ood_activity(sequence) or \ EgoExo4D_val_ood_pair(sequence) def EgoExo4D_val_id_capture(sequence): # NOTE(bvh): this still achieves no TAKE level overlap, as desired is_ood = EgoExo4D_val_ood(sequence) heldout_capture = sequence.split('_')[-2][-1] == '2' return not(is_ood) and heldout_capture if dataset == 'EgoExo4D': save_split('train5', EgoExo4D_train) save_split('val5_ood_site', EgoExo4D_val_ood_site) save_split('val5_ood_activity', EgoExo4D_val_ood_activity) save_split('val5_ood_pair', EgoExo4D_val_ood_pair) save_split('val5_ood', EgoExo4D_val_ood) save_split('val5_id_capture', EgoExo4D_val_id_capture) ####################################### Kubric4D # Official GCD train/test/val split def Kubric4D_train(sequence): return int(sequence[3:]) < 2800 def Kubric4D_val(sequence): return 2800 <= int(sequence[3:]) < 2900 def Kubric4D_test(sequence): return int(sequence[3:]) >= 2900 if dataset == 'Kubric4D': save_split('train', Kubric4D_train) save_split('val', Kubric4D_val) save_split('test', Kubric4D_test) ####################################### Kubric5D # Proposead train/val/test split (similar to Kubric4D) def Kubric5D_train(sequence): return int(sequence[3:]) < 9600 def Kubric5D_val(sequence): return 9600 <= int(sequence[3:]) < 9800 def Kubric5D_test(sequence): return int(sequence[3:]) >= 9800 if dataset == 'Kubric5D': save_split('train', Kubric5D_train) save_split('val', Kubric5D_val) save_split('test', Kubric5D_test) ####################################### LBMv12 # Sequences with the rightmost digit == 4 are val, rest is train def LBMv12_train(sequence): return sequence.split('__')[-1][-1] != '4' def LBMv12_val(sequence): return sequence.split('__')[-1][-1] == '4' if dataset == 'LBMv12': save_split('train', LBMv12_train) save_split('val', LBMv12_val) ####################################### LyftL5 # Sequences starting with 'train' are train, with 'test' are test def LyftL5_train(sequence): return sequence.startswith('train') def LyftL5_test(sequence): return sequence.startswith('test') if dataset == 'LyftL5': save_split('train', LyftL5_train) save_split('test', LyftL5_test) ####################################### MVImgNet # Sequences with the rightmost digit == 5 are val, rest is train def MVImgNet_train(sequence): return sequence.split('__')[-1][-1] != '5' def MVImgNet_val(sequence): return sequence.split('__')[-1][-1] == '5' if dataset == 'MVImgNet': save_split('train', MVImgNet_train) save_split('val', MVImgNet_val) ####################################### PD4D # Official GCD split def PD4D_train(sequence): split = sequence.split('_')[-1] return split not in PD4Dval and split not in PD4Dtest def PD4D_val(sequence): split = sequence.split('_')[-1] return split in PD4Dval def PD4D_test(sequence): split = sequence.split('_')[-1] return split in PD4Dtest if dataset == 'PD4D': save_split('train', PD4D_train) save_split('val', PD4D_val) save_split('test', PD4D_test) ####################################### RealEstate10k # Sequences labeled with "train" are train, rest is val def RealEstate10k_train(sequence): return sequence.split('__')[0] == 'train' def RealEstate10k_val(sequence): return sequence.split('__')[0] == 'test' if dataset == 'RealEstate10k': save_split('train', RealEstate10k_train) save_split('val', RealEstate10k_val) ####################################### ScanNetVideo # Sequences with the rightmost digit == 0 are val, rest is train def ScanNetVideo_train(sequence): return sequence.split('__')[1].split('_')[0][-1] != '0' def ScanNetVideo_val(sequence): return sequence.split('__')[1].split('_')[0][-1] == '0' if dataset == 'ScanNetVideo': save_split('train', ScanNetVideo_train) save_split('val', ScanNetVideo_val) ####################################### LyftL5 # Sequences starting with 'train' are train, with 'val' are val def Waymo_train(sequence): return sequence.startswith('train') def Waymo_val(sequence): return sequence.startswith('val') if dataset == 'Waymo': save_split('train', Waymo_train) save_split('val', Waymo_val) ####################################### WildRGBD # Sequences numbered at 4 are val, rest is train def WildRGBD_train(sequence): return int(sequence.split('_')[-1]) != 4 def WildRGBD_val(sequence): return int(sequence.split('_')[-1]) == 4 if dataset == 'WildRGBD': save_split('train', WildRGBD_train) save_split('val', WildRGBD_val) ####################################### if __name__ == "__main__": path = sys.argv[1] if path.endswith('/'): path = path[:-1] folders = list_s3(path) if len(sys.argv) == 3: nested = int(sys.argv[2]) else: nested = 0 if 'tarfiles' in folders: create_indices_single(f'{path}/tarfiles', nested) else: for f1 in folders: if f1 == 'filenames': continue folders2 = list_s3(f'{path}/{f1}') if 'tarfiles' in folders2: create_indices_single(f'{path}/{f1}/tarfiles', nested) else: for f2 in folders2: if f2 == 'filenames': continue if not f2.startswith('filenames'): create_indices_single(f'{path}/{f1}/{f2}/tarfiles', nested)