Skip to content

Commit ba669c1

Browse files
author
Jan Czarnowski
committed
Progress bar; moved samples to samples folder
1 parent 155f2d6 commit ba669c1

9 files changed

+883
-18
lines changed

.gitmodules

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
[submodule "ScanNet"]
2+
path = ScanNet
3+
url = https://github.com/ScanNet/ScanNet

ScanNet

Submodule ScanNet added at b8c2a05

download-scannet.py

+204
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
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()

download_sample.sh

+66-18
Original file line numberDiff line numberDiff line change
@@ -1,37 +1,85 @@
11
#!/bin/bash
22

3-
SENS_PROG=/home/jancz/dev/co/ScanNet/SensReader/c++/sens
4-
OUT_DIR=$2
3+
reader_dir=ScanNet/SensReader/c++
4+
reader_exec="$reader_dir/sens"
55

6-
if [ "$#" -ne 2 ]; then
7-
echo "Usage: $0 <sample_file> out_dir"
8-
exit
9-
fi
6+
# Progress bar printing snagged from
7+
# https://github.com/fearside/ProgressBar
8+
function progress_bar {
9+
let _progress=(${1}*100/${2}*100)/100
10+
let _done=(${_progress}*4)/10
11+
let _left=40-$_done
12+
_fill=$(printf "%${_done}s")
13+
_empty=$(printf "%${_left}s")
14+
printf "\rProgress : [${_fill// /\#}${_empty// /-}] ${_progress}%%"
15+
}
1016

1117
# Downloads and uncompresses the .sens file given a scene id
1218
# e.g. 'scene0425_00'
1319
function download_scene_images {
14-
scene_id=$1
20+
_scene_id=$1
21+
_out_dir=$2
1522

16-
python download-scannet.py -o . --id $scene_id --type .sens || exit
23+
python download-scannet.py -o . --id ${_scene_id} --type .sens &> /dev/null
1724

1825
# decompress .sens file
19-
sens_file="scans/$line/$line.sens"
20-
$SENS_PROG $sens_file "$OUT_DIR/$line" || exit
21-
22-
# remove .sens file
23-
rm -rf "$sens_file"
26+
sens_file="scans/${_scene_id}/${_scene_id}.sens"
27+
${reader_exec} ${sens_file} "${_out_dir}/${_scene_id}" &> /dev/null
2428
}
2529

30+
# Downloads model ply file for a given scene 'scene0425_00'
2631
function download_scene_ply {
27-
scene_id=$1
28-
python download-scannet.py -o $OUT_DIR/$scene_id --id $scene_id --type _vh_clean_2.ply || exit
32+
_scene_id=$1
33+
_out_dir=$2
34+
python download-scannet.py -o . --id ${_scene_id} --type _vh_clean_2.ply &> /dev/null
35+
36+
# move the downloaded ply into proper scene directory
37+
mv "scans/${_scene_id}/${_scene_id}_vh_clean_2.ply" "${_out_dir}/${_scene_id}/"
38+
}
39+
40+
# Ensure the reader is compiled
41+
function compile_reader {
42+
if [[ ! -f $reader_exec ]]; then
43+
echo "Compiling the sens reader"
44+
make -C $reader_dir
45+
fi
2946
}
3047

48+
# check command line args
49+
if [ "$#" -ne 2 ]; then
50+
echo "Usage: $0 <sample_file> <out_dir>"
51+
exit
52+
fi
53+
54+
# parse args
55+
sample_file=$1
56+
out_dir=$2
57+
58+
# validate args
59+
if [[ ! -f $sample_file ]]; then
60+
echo "Sample file ${sample_file} does not exist"
61+
exit -1
62+
fi
63+
64+
# ensure that sens reader is compiled
65+
compile_reader || exit
66+
67+
# probe the input file for number of scenes
68+
num_scenes=$(cat $1 | wc -l)
69+
70+
# create the output folder
71+
mkdir -p ${out_dir} || exit
72+
73+
# download the data
74+
echo "Downloading data for $num_scenes scenes"
75+
i=0
3176
while read scene; do
32-
echo "Downloading data for scene $scene"
77+
progress_bar ${i} ${num_scenes}
78+
download_scene_images ${scene} ${out_dir} || exit
79+
download_scene_ply ${scene} ${out_dir} || exit
3380

34-
#download_scene_images $scene
35-
download_scene_ply $scene
81+
# clear the downloaded files
82+
rm -rf scans
3683

84+
i=$((i + 1))
3785
done < $1

0 commit comments

Comments
 (0)