|
| 1 | +#!/usr/bin/env python |
| 2 | +# Downloads ScanNet public data release |
| 3 | +# Run with ./download-scannet.py (or python download-scannet.py on Windows) |
| 4 | +# -*- coding: utf-8 -*- |
| 5 | +import argparse |
| 6 | +import os |
| 7 | +import urllib.request as urllib #(for python3) |
| 8 | +#import urllib |
| 9 | +import tempfile |
| 10 | + |
| 11 | +BASE_URL = 'http://dovahkiin.stanford.edu/scannet-public/' |
| 12 | +TOS_URL = BASE_URL + 'ScanNet_TOS.pdf' |
| 13 | +FILETYPES = ['.aggregation.json', '.sens', '.txt', '_vh_clean.ply', '_vh_clean_2.0.010000.segs.json', '_vh_clean_2.ply', '_vh_clean.segs.json', '_vh_clean.aggregation.json', '_vh_clean_2.labels.ply', '_2d-instance.zip', '_2d-instance-filt.zip', '_2d-label.zip', '_2d-label-filt.zip'] |
| 14 | +FILETYPES_TEST = ['.sens', '.txt', '_vh_clean.ply', '_vh_clean_2.ply'] |
| 15 | +PREPROCESSED_FRAMES_FILE = ['scannet_frames_25k.zip', '5.6GB'] |
| 16 | +TEST_FRAMES_FILE = ['scannet_frames_test.zip', '610MB'] |
| 17 | +LABEL_MAP_FILES = ['scannetv2-labels.combined.tsv', 'scannet-labels.combined.tsv'] |
| 18 | +RELEASES = ['v2/scans', 'v1/scans'] |
| 19 | +RELEASES_TASKS = ['v2/tasks', 'v1/tasks'] |
| 20 | +RELEASES_NAMES = ['v2', 'v1'] |
| 21 | +RELEASE = RELEASES[0] |
| 22 | +RELEASE_TASKS = RELEASES_TASKS[0] |
| 23 | +RELEASE_NAME = RELEASES_NAMES[0] |
| 24 | +LABEL_MAP_FILE = LABEL_MAP_FILES[0] |
| 25 | +RELEASE_SIZE = '1.2TB' |
| 26 | +V1_IDX = 1 |
| 27 | + |
| 28 | + |
| 29 | +def get_release_scans(release_file): |
| 30 | + #scan_lines = urllib.request.urlopen(release_file) |
| 31 | + scan_lines = urllib.urlopen(release_file) |
| 32 | + scans = [] |
| 33 | + for scan_line in scan_lines: |
| 34 | + scan_id = scan_line.decode('utf8').rstrip('\n') |
| 35 | + scans.append(scan_id) |
| 36 | + return scans |
| 37 | + |
| 38 | + |
| 39 | +def download_release(release_scans, out_dir, file_types, use_v1_sens): |
| 40 | + if len(release_scans) == 0: |
| 41 | + return |
| 42 | + print('Downloading ScanNet ' + RELEASE_NAME + ' release to ' + out_dir + '...') |
| 43 | + for scan_id in release_scans: |
| 44 | + scan_out_dir = os.path.join(out_dir, scan_id) |
| 45 | + download_scan(scan_id, scan_out_dir, file_types, use_v1_sens) |
| 46 | + print('Downloaded ScanNet ' + RELEASE_NAME + ' release.') |
| 47 | + |
| 48 | + |
| 49 | +def download_file(url, out_file): |
| 50 | + out_dir = os.path.dirname(out_file) |
| 51 | + if not os.path.isdir(out_dir): |
| 52 | + os.makedirs(out_dir) |
| 53 | + if not os.path.isfile(out_file): |
| 54 | + print('\t' + url + ' > ' + out_file) |
| 55 | + fh, out_file_tmp = tempfile.mkstemp(dir=out_dir) |
| 56 | + f = os.fdopen(fh, 'w') |
| 57 | + f.close() |
| 58 | + #urllib.request.urlretrieve(url, out_file_tmp) |
| 59 | + urllib.urlretrieve(url, out_file_tmp) |
| 60 | + os.rename(out_file_tmp, out_file) |
| 61 | + else: |
| 62 | + print('WARNING: skipping download of existing file ' + out_file) |
| 63 | + |
| 64 | +def download_scan(scan_id, out_dir, file_types, use_v1_sens): |
| 65 | + print('Downloading ScanNet ' + RELEASE_NAME + ' scan ' + scan_id + ' ...') |
| 66 | + if not os.path.isdir(out_dir): |
| 67 | + os.makedirs(out_dir) |
| 68 | + for ft in file_types: |
| 69 | + v1_sens = use_v1_sens and ft == '.sens' |
| 70 | + url = BASE_URL + RELEASE + '/' + scan_id + '/' + scan_id + ft if not v1_sens else BASE_URL + RELEASES[V1_IDX] + '/' + scan_id + '/' + scan_id + ft |
| 71 | + out_file = out_dir + '/' + scan_id + ft |
| 72 | + download_file(url, out_file) |
| 73 | + print('Downloaded scan ' + scan_id) |
| 74 | + |
| 75 | + |
| 76 | +def download_task_data(out_dir): |
| 77 | + print('Downloading ScanNet v1 task data...') |
| 78 | + files = [ |
| 79 | + LABEL_MAP_FILES[V1_IDX], 'obj_classification/data.zip', |
| 80 | + 'obj_classification/trained_models.zip', 'voxel_labeling/data.zip', |
| 81 | + 'voxel_labeling/trained_models.zip' |
| 82 | + ] |
| 83 | + for file in files: |
| 84 | + url = BASE_URL + RELEASES_TASKS[V1_IDX] + '/' + file |
| 85 | + localpath = os.path.join(out_dir, file) |
| 86 | + localdir = os.path.dirname(localpath) |
| 87 | + if not os.path.isdir(localdir): |
| 88 | + os.makedirs(localdir) |
| 89 | + download_file(url, localpath) |
| 90 | + print('Downloaded task data.') |
| 91 | + |
| 92 | + |
| 93 | +def download_label_map(out_dir): |
| 94 | + print('Downloading ScanNet ' + RELEASE_NAME + ' label mapping file...') |
| 95 | + files = [ LABEL_MAP_FILE ] |
| 96 | + for file in files: |
| 97 | + url = BASE_URL + RELEASE_TASKS + '/' + file |
| 98 | + localpath = os.path.join(out_dir, file) |
| 99 | + localdir = os.path.dirname(localpath) |
| 100 | + if not os.path.isdir(localdir): |
| 101 | + os.makedirs(localdir) |
| 102 | + download_file(url, localpath) |
| 103 | + print('Downloaded ScanNet ' + RELEASE_NAME + ' label mapping file.') |
| 104 | + |
| 105 | + |
| 106 | +def main(): |
| 107 | + parser = argparse.ArgumentParser(description='Downloads ScanNet public data release.') |
| 108 | + parser.add_argument('-o', '--out_dir', required=True, help='directory in which to download') |
| 109 | + parser.add_argument('--task_data', action='store_true', help='download task data (v1)') |
| 110 | + parser.add_argument('--label_map', action='store_true', help='download label map file') |
| 111 | + parser.add_argument('--v1', action='store_true', help='download ScanNet v1 instead of v2') |
| 112 | + parser.add_argument('--id', help='specific scan id to download') |
| 113 | + parser.add_argument('--preprocessed_frames', action='store_true', help='download preprocessed subset of ScanNet frames (' + PREPROCESSED_FRAMES_FILE[1] + ')') |
| 114 | + parser.add_argument('--test_frames_2d', action='store_true', help='download 2D test frames (' + TEST_FRAMES_FILE[1] + '; also included with whole dataset download)') |
| 115 | + parser.add_argument('--type', help='specific file type to download (.aggregation.json, .sens, .txt, _vh_clean.ply, _vh_clean_2.0.010000.segs.json, _vh_clean_2.ply, _vh_clean.segs.json, _vh_clean.aggregation.json, _vh_clean_2.labels.ply, _2d-instance.zip, _2d-instance-filt.zip, _2d-label.zip, _2d-label-filt.zip)') |
| 116 | + args = parser.parse_args() |
| 117 | + |
| 118 | + #print('By pressing any key to continue you confirm that you have agreed to the ScanNet terms of use as described at:') |
| 119 | + #print(TOS_URL) |
| 120 | + #print('***') |
| 121 | + #print('Press any key to continue, or CTRL-C to exit.') |
| 122 | + #key = input('') |
| 123 | + |
| 124 | + if args.v1: |
| 125 | + global RELEASE |
| 126 | + global RELEASE_TASKS |
| 127 | + global RELEASE_NAME |
| 128 | + global LABEL_MAP_FILE |
| 129 | + RELEASE = RELEASES[V1_IDX] |
| 130 | + RELEASE_TASKS = RELEASES_TASKS[V1_IDX] |
| 131 | + RELEASE_NAME = RELEASES_NAMES[V1_IDX] |
| 132 | + LABEL_MAP_FILE = LABEL_MAP_FILES[V1_IDX] |
| 133 | + |
| 134 | + release_file = BASE_URL + RELEASE + '.txt' |
| 135 | + release_scans = get_release_scans(release_file) |
| 136 | + file_types = FILETYPES; |
| 137 | + release_test_file = BASE_URL + RELEASE + '_test.txt' |
| 138 | + release_test_scans = get_release_scans(release_test_file) |
| 139 | + file_types_test = FILETYPES_TEST; |
| 140 | + out_dir_scans = os.path.join(args.out_dir, 'scans') |
| 141 | + out_dir_test_scans = os.path.join(args.out_dir, 'scans_test') |
| 142 | + out_dir_tasks = os.path.join(args.out_dir, 'tasks') |
| 143 | + |
| 144 | + if args.type: # download file type |
| 145 | + file_type = args.type |
| 146 | + if file_type not in FILETYPES: |
| 147 | + print('ERROR: Invalid file type: ' + file_type) |
| 148 | + return |
| 149 | + file_types = [file_type] |
| 150 | + if file_type in FILETYPES_TEST: |
| 151 | + file_types_test = [file_type] |
| 152 | + else: |
| 153 | + file_types_test = [] |
| 154 | + if args.task_data: # download task data |
| 155 | + download_task_data(out_dir_tasks) |
| 156 | + elif args.label_map: # download label map file |
| 157 | + download_label_map(args.out_dir) |
| 158 | + elif args.preprocessed_frames: # download preprocessed scannet_frames_25k.zip file |
| 159 | + if args.v1: |
| 160 | + print('ERROR: Preprocessed frames only available for ScanNet v2') |
| 161 | + print('You are downloading the preprocessed subset of frames ' + PREPROCESSED_FRAMES_FILE[0] + ' which requires ' + PREPROCESSED_FRAMES_FILE[1] + ' of space.') |
| 162 | + download_file(os.path.join(BASE_URL, RELEASE_TASKS, PREPROCESSED_FRAMES_FILE[0]), os.path.join(out_dir_tasks, PREPROCESSED_FRAMES_FILE[0])) |
| 163 | + elif args.test_frames_2d: # download test scannet_frames_test.zip file |
| 164 | + if args.v1: |
| 165 | + print('ERROR: 2D test frames only available for ScanNet v2') |
| 166 | + print('You are downloading the 2D test set ' + TEST_FRAMES_FILE[0] + ' which requires ' + TEST_FRAMES_FILE[1] + ' of space.') |
| 167 | + download_file(os.path.join(BASE_URL, RELEASE_TASKS, TEST_FRAMES_FILE[0]), os.path.join(out_dir_tasks, TEST_FRAMES_FILE[0])) |
| 168 | + elif args.id: # download single scan |
| 169 | + scan_id = args.id |
| 170 | + is_test_scan = scan_id in release_test_scans |
| 171 | + if scan_id not in release_scans and (not is_test_scan or args.v1): |
| 172 | + print('ERROR: Invalid scan id: ' + scan_id) |
| 173 | + else: |
| 174 | + out_dir = os.path.join(out_dir_scans, scan_id) if not is_test_scan else os.path.join(out_dir_test_scans, scan_id) |
| 175 | + scan_file_types = file_types if not is_test_scan else file_types_test |
| 176 | + use_v1_sens = not is_test_scan |
| 177 | + #if not is_test_scan and not args.v1 and '.sens' in scan_file_types: |
| 178 | + #print('Note: ScanNet v2 uses the same .sens files as ScanNet v1: Press \'n\' to exclude downloading .sens files for each scan') |
| 179 | + #key = input('') |
| 180 | + #if key.strip().lower() == 'n': |
| 181 | + # scan_file_types.remove('.sens') |
| 182 | + download_scan(scan_id, out_dir, scan_file_types, use_v1_sens) |
| 183 | + else: # download entire release |
| 184 | + if len(file_types) == len(FILETYPES): |
| 185 | + print('WARNING: You are downloading the entire ScanNet ' + RELEASE_NAME + ' release which requires ' + RELEASE_SIZE + ' of space.') |
| 186 | + else: |
| 187 | + print('WARNING: You are downloading all ScanNet ' + RELEASE_NAME + ' scans of type ' + file_types[0]) |
| 188 | + print('Note that existing scan directories will be skipped. Delete partially downloaded directories to re-download.') |
| 189 | + print('***') |
| 190 | + print('Press any key to continue, or CTRL-C to exit.') |
| 191 | + key = input('') |
| 192 | + if not args.v1 and '.sens' in file_types: |
| 193 | + print('Note: ScanNet v2 uses the same .sens files as ScanNet v1: Press \'n\' to exclude downloading .sens files for each scan') |
| 194 | + key = input('') |
| 195 | + if key.strip().lower() == 'n': |
| 196 | + file_types.remove('.sens') |
| 197 | + download_release(release_scans, out_dir_scans, file_types, use_v1_sens=True) |
| 198 | + if not args.v1: |
| 199 | + download_label_map(args.out_dir) |
| 200 | + download_release(release_test_scans, out_dir_test_scans, file_types_test, use_v1_sens=False) |
| 201 | + download_file(os.path.join(BASE_URL, RELEASE_TASKS, TEST_FRAMES_FILE[0]), os.path.join(out_dir_tasks, TEST_FRAMES_FILE[0])) |
| 202 | + |
| 203 | + |
| 204 | +if __name__ == "__main__": main() |
0 commit comments