Skip to content

Commit 7fad5b1

Browse files
hychiang-githychiang-git
committed
done collect_point
1 parent 5f94952 commit 7fad5b1

File tree

6 files changed

+1501
-104
lines changed

6 files changed

+1501
-104
lines changed

‎prepare_data/collect_points/README.md

+108
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# Extract scene points from ply files in ScanNet v2 dataset
2+
3+
The codes are borrowed from [PointNet++](https://github.com/charlesq34/pointnet2/tree/master/scannet/preprocessing)
4+
5+
6+
7+
## Example usage
8+
9+
```
10+
// training set
11+
python3 trainval_point.py
12+
--scannet_path /path/to/scannet/scans/
13+
--label_map_file /path/to/scannetv2-labels.combined.tsv
14+
--num_proc 5
15+
16+
17+
// testing set
18+
python3 test_point.py \
19+
--scannet_path /path/to/scannet/scans_test/ \
20+
--num_proc 5 \
21+
```
22+
23+
<!--
24+
## Getting Started
25+
26+
These instructions will get you a copy of the project up and running on your local machine for development and testing purposes. See deployment for notes on how to deploy the project on a live system.
27+
28+
### Prerequisites
29+
30+
What things you need to install the software and how to install them
31+
32+
```
33+
Give examples
34+
```
35+
36+
### Installing
37+
38+
A step by step series of examples that tell you how to get a development env running
39+
40+
Say what the step will be
41+
42+
```
43+
Give the example
44+
```
45+
46+
And repeat
47+
48+
```
49+
until finished
50+
```
51+
52+
End with an example of getting some data out of the system or using it for a little demo
53+
54+
## Running the tests
55+
56+
Explain how to run the automated tests for this system
57+
58+
### Break down into end to end tests
59+
60+
Explain what these tests test and why
61+
62+
```
63+
Give an example
64+
```
65+
66+
### And coding style tests
67+
68+
Explain what these tests test and why
69+
70+
```
71+
Give an example
72+
```
73+
74+
## Deployment
75+
76+
Add additional notes about how to deploy this on a live system
77+
78+
## Built With
79+
80+
* [Dropwizard](http://www.dropwizard.io/1.0.2/docs/) - The web framework used
81+
* [Maven](https://maven.apache.org/) - Dependency Management
82+
* [ROME](https://rometools.github.io/rome/) - Used to generate RSS Feeds
83+
84+
## Contributing
85+
86+
Please read [CONTRIBUTING.md](https://gist.github.com/PurpleBooth/b24679402957c63ec426) for details on our code of conduct, and the process for submitting pull requests to us.
87+
88+
## Versioning
89+
90+
We use [SemVer](http://semver.org/) for versioning. For the versions available, see the [tags on this repository](https://github.com/your/project/tags).
91+
92+
## Authors
93+
94+
* **Billie Thompson** - *Initial work* - [PurpleBooth](https://github.com/PurpleBooth)
95+
96+
See also the list of [contributors](https://github.com/your/project/contributors) who participated in this project.
97+
98+
## License
99+
100+
This project is licensed under the MIT License - see the [LICENSE.md](LICENSE.md) file for details
101+
102+
## Acknowledgments
103+
104+
* Hat tip to anyone whose code was used
105+
* Inspiration
106+
* etc
107+
--->
108+

‎prepare_data/collect_points/test_point.py

+31-53
Original file line numberDiff line numberDiff line change
@@ -6,71 +6,49 @@
66
import numpy as np
77
import multiprocessing as mp
88
from functools import partial
9-
from utils import pc_util
10-
from utils import scannet_util
11-
g_label_names = scannet_util.g_label_names
12-
g_label_ids = scannet_util.g_label_ids
139

14-
SCANNET_DIR = None # '/tmp3/hychiang/ScanNet.v2/ScanNet/scans/'
15-
SCENE_NAMES = None # [line.rstrip() for line in open('./Benchmark/scannetv2_train.txt')]
16-
OUTPUT_FOLDER = None # 'scans_train'
17-
LOG_FILE = 'log.txt'
18-
LOG_FOUT = None
10+
sys.path.append('../utils')
11+
import pc_utils
12+
import scannet_utils
1913

20-
def collect_one_scene_data(scene_name, out_filename):
14+
'''
15+
params
16+
'''
17+
parser = argparse.ArgumentParser()
18+
parser.add_argument('--scannet_path', required=True, help='path to scannet data')
19+
parser.add_argument("--num_proc", required=False, type=int, default=30, help="number of parallel process, default is 30")
20+
opt = parser.parse_args()
21+
22+
23+
def collect_point_data(scene_name):
2124
# Over-segmented segments: maps from segment to vertex/point IDs
22-
data_folder = os.path.join(SCANNET_DIR, scene_name)
25+
data_folder = os.path.join(opt.scannet_path, scene_name)
2326
ply_filename = os.path.join(data_folder, '%s_vh_clean_2.ply' % (scene_name))
24-
points = pc_util.read_ply_rgba_normal(ply_filename)
27+
points = pc_utils.read_ply_rgba_normal(ply_filename)
2528

2629
points = np.delete(points, 6, 1) # only RGB, ignoring A
2730
data = points
31+
out_filename = os.path.join(data_folder, scene_name+'.npy') # scene0000_00/scene0000_00.npy
2832
np.save(out_filename, data)
29-
log_string(scene_name+' save to '+out_filename+', point:'+str(points.shape)+', data:'+str(data.shape))
33+
print(scene_name, ' points shape:', data.shape)
3034

31-
def log_string(out_str):
32-
LOG_FOUT.write(out_str+'\n')
33-
LOG_FOUT.flush()
34-
print(out_str)
3535

3636
def preprocess_scenes(scene_name):
37-
log_string(scene_name)
3837
try:
39-
out_dir = os.path.join(OUTPUT_FOLDER, scene_name) # scene0000_00/scene0000_00.npy
40-
if not os.path.exists(out_dir):
41-
os.mkdir(out_dir)
42-
out_file = os.path.join(out_dir, scene_name+'.npy') # scene0000_00/scene0000_00.npy
43-
collect_one_scene_data(scene_name, out_file)
38+
collect_point_data(scene_name)
4439
except Exception as e:
45-
log_string(scene_name+'ERROR!!')
46-
log_string(str(e))
47-
48-
if __name__=='__main__':
49-
parser = argparse.ArgumentParser()
50-
parser.add_argument("--scene_list", required=True, help="scannet split scene list, e.g. ./Benchmark/scannetv2_train.txt")
51-
parser.add_argument("--scannet_dir",required=True, help="scannet data dir, e.g. {path/to/scannet/data/dir}/scans or {path/to/scannet/data/dir}/scans_test")
52-
parser.add_argument("--output_dir", required=True, help="output dir (folder), e.g. ./scans_train")
53-
parser.add_argument("--num_proc", required=False, type=int, default=30, help="number of parallel process, default is 30")
54-
args = parser.parse_args()
55-
SCENE_NAMES = [line.rstrip() for line in open(args.scene_list)]
56-
SCANNET_DIR = args.scannet_dir
57-
OUTPUT_FOLDER = args.output_dir
40+
sys.stderr.write(scene_name+'ERROR!!')
41+
sys.stderr.write(str(e))
42+
sys.exit(-1)
43+
44+
def main():
45+
scenes = [d for d in os.listdir(opt.scannet_path) if os.path.isdir(os.path.join(opt.scannet_path, d))]
46+
scenes.sort()
47+
print('Find %d scenes' % len(scenes))
48+
print('Extract points (Vertex XYZ, RGB, NxNyNx)')
5849

59-
print('*** Total Scene in list: ', len(SCENE_NAMES))
60-
print('*** ScanNet Data Directory: ', SCANNET_DIR)
61-
print('*** Output Directory: ', OUTPUT_FOLDER)
62-
print('*** NUM of Processes to parallel: ', args.num_proc)
63-
print('*** Extract points (Vertex XYZ, RGB, NxNyNx, Label, Instance-label) parallel in 5 Seconds***')
64-
65-
if not os.path.exists(OUTPUT_FOLDER):
66-
os.mkdir(OUTPUT_FOLDER)
67-
LOG_FOUT = open(os.path.join(OUTPUT_FOLDER, LOG_FILE),'w')
68-
time.sleep(5)
50+
pool = mp.Pool(opt.num_proc)
51+
pool.map(preprocess_scenes, scenes)
6952

70-
print('*** GO ***')
71-
72-
73-
pool = mp.Pool(args.num_proc)
74-
pool.map(preprocess_scenes, SCENE_NAMES)
75-
76-
LOG_FOUT.close()
53+
if __name__=='__main__':
54+
main()

‎prepare_data/collect_points/trainval_point.py

+28-51
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,22 @@
1313
g_label_names = scannet_utils.g_label_names
1414
g_label_ids = scannet_utils.g_label_ids
1515

16-
SCANNET_DIR = None # '/tmp3/hychiang/ScanNet.v2/ScanNet/scans/'
17-
SCENE_NAMES = None # [line.rstrip() for line in open('./Benchmark/scannetv2_train.txt')]
18-
LABEL_MAP_FILE = None # './scannetv2-labels.combined.tsv'
19-
OUTPUT_FOLDER = None # 'scans_train'
20-
LOG_FILE = 'log.txt'
21-
LOG_FOUT = None
16+
'''
17+
params
18+
'''
19+
parser = argparse.ArgumentParser()
20+
parser.add_argument('--scannet_path', required=True, help='path to scannet data')
21+
parser.add_argument('--label_map_file', default='', help='path to scannetv2-labels.combined.tsv (required for label export only)')
22+
parser.add_argument("--num_proc", required=False, type=int, default=30, help="number of parallel process, default is 30")
23+
opt = parser.parse_args()
2224

23-
def collect_one_scene_data_label(scene_name, out_filename):
25+
26+
def collect_point_data(scene_name):
2427
# read label mapping file
25-
label_map = scannet_utils.read_label_mapping(LABEL_MAP_FILE, label_from='raw_category', label_to='nyu40id')
28+
label_map = scannet_utils.read_label_mapping(opt.label_map_file, label_from='raw_category', label_to='nyu40id')
2629

2730
# Over-segmented segments: maps from segment to vertex/point IDs
28-
data_folder = os.path.join(SCANNET_DIR, scene_name)
31+
data_folder = os.path.join(opt.scannet_path, scene_name)
2932
# Read segmentation label
3033
seg_filename = os.path.join(data_folder, '%s_vh_clean_2.0.010000.segs.json'%(scene_name))
3134
seg_to_verts, num_verts = scannet_utils.read_segmentation(seg_filename)
@@ -69,55 +72,29 @@ def collect_one_scene_data_label(scene_name, out_filename):
6972
#print(data.shape)
7073
#for i in range(20):
7174
# print(data[i, 10])
75+
out_filename = os.path.join(data_folder, scene_name+'.npy') # scene0000_00/scene0000_00.npy
7276
np.save(out_filename, data)
73-
log_string(scene_name+' save to '+out_filename+', with data: point:'+str(points.shape)+', label:'+str(label_ids.shape)+', instance label:'+str(instance_ids.shape)+', data:'+str(data.shape))
7477

75-
def log_string(out_str):
76-
LOG_FOUT.write(out_str+'\n')
77-
LOG_FOUT.flush()
78-
print(out_str)
78+
print(scene_name, ' points shape:', data.shape)
79+
7980

8081
def preprocess_scenes(scene_name):
81-
log_string(scene_name)
8282
try:
83-
out_dir = os.path.join(OUTPUT_FOLDER, scene_name) # scene0000_00/scene0000_00.npy
84-
if not os.path.exists(out_dir):
85-
os.mkdir(out_dir)
86-
out_file = os.path.join(out_dir, scene_name+'.npy') # scene0000_00/scene0000_00.npy
87-
collect_one_scene_data_label(scene_name, out_file)
83+
collect_point_data(scene_name)
8884
except Exception as e:
89-
log_string(scene_name+'ERROR!!')
90-
log_string(str(e))
85+
sys.stderr.write(scene_name+'ERROR!!')
86+
sys.stderr.write(str(e))
87+
sys.exit(-1)
9188

92-
if __name__=='__main__':
93-
parser = argparse.ArgumentParser()
94-
parser.add_argument("--scene_list", required=True, help="scannet split scene list, e.g. ./Benchmark/scannetv2_train.txt")
95-
parser.add_argument("--label_map_file", required=True, help="scannet label mapping file , e.g. ./scannetv2-labels.combined.tsv")
96-
parser.add_argument("--scannet_dir",required=True, help="scannet data dir, e.g. {path/to/scannet/data/dir}/scans or {path/to/scannet/data/dir}/scans_test")
97-
parser.add_argument("--output_dir", required=True, help="output dir (folder), e.g. ./scans_train")
98-
parser.add_argument("--num_proc", required=False, type=int, default=30, help="number of parallel process, default is 30")
99-
args = parser.parse_args()
100-
LABEL_MAP_FILE = args.label_map_file
101-
SCENE_NAMES = [line.rstrip() for line in open(args.scene_list)]
102-
SCANNET_DIR = args.scannet_dir
103-
OUTPUT_FOLDER = args.output_dir
104-
105-
print('*** Total Scene in list: ', len(SCENE_NAMES))
106-
print('*** Read Label Mapping File from: ', LABEL_MAP_FILE)
107-
print('*** ScanNet Data Directory: ', SCANNET_DIR)
108-
print('*** Output Directory: ', OUTPUT_FOLDER)
109-
print('*** NUM of Processes to parallel: ', args.num_proc)
110-
print('*** Extract points (Vertex XYZ, RGB, NxNyNx, Label, Instance-label) parallel in 5 Seconds***')
111-
112-
if not os.path.exists(OUTPUT_FOLDER):
113-
os.mkdir(OUTPUT_FOLDER)
114-
LOG_FOUT = open(os.path.join(OUTPUT_FOLDER, LOG_FILE),'w')
115-
time.sleep(5)
116-
117-
print('*** GO ***')
11889

90+
def main():
91+
scenes = [d for d in os.listdir(opt.scannet_path) if os.path.isdir(os.path.join(opt.scannet_path, d))]
92+
scenes.sort()
93+
print('Find %d scenes' % len(scenes))
94+
print('Extract points (Vertex XYZ, RGB, NxNyNx, Label, Instance-label)')
11995

120-
pool = mp.Pool(args.num_proc)
121-
pool.map(preprocess_scenes, SCENE_NAMES)
96+
pool = mp.Pool(opt.num_proc)
97+
pool.map(preprocess_scenes, scenes)
12298

123-
LOG_FOUT.close()
99+
if __name__=='__main__':
100+
main()

0 commit comments

Comments
 (0)