-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmake_json_files.py
61 lines (53 loc) · 2.56 KB
/
make_json_files.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import os
import numpy as np
import glob
import random
import json
def get_train_test_dict(root, train_r):
"""
Make json files of images and masks paths.
Format is {0: {x: path_to_image, m: path_to_mask}, 1: ...}
train_r: ratio of train test
"""
filenames = [os.path.basename(elem) for elem in glob.glob(os.path.join(root, "Normal/*"))]
random.shuffle(filenames)
train_index = int(train_r*len(filenames))
train = filenames[:train_index]
test = filenames[train_index:]
train_dict = []
for i in range(len(train)):
train_dict.append( {'x': os.path.join(root, "Normal",train[i]), 'm': os.path.join(root, "Mask",train[i])})
test_dict = []
for i in range(len(test)):
test_dict.append( {'x': os.path.join(root, "Normal",test[i]), 'm': os.path.join(root, "Mask",test[i])})
return(train_dict, test_dict)
def get_train_test_dict_depth(root, train_r):
"""
Make json files of images and masksand depth paths.
Format is {0: {x: path_to_image, m: path_to_mask}, 1: ...}
train_r: ratio of train test
"""
filenames = [os.path.basename(elem) for elem in glob.glob(os.path.join(root, "Normal/*"))]
random.shuffle(filenames)
train_index = int(train_r*len(filenames))
train = filenames[:train_index]
test = filenames[train_index:]
train_dict = []
for i in range(len(train)):
train_dict.append( {'x': os.path.join(root, "Normal",train[i]), 'm': os.path.join(root, "Mask",train[i]), 'd': os.path.join(root, "Depth",train[i])})
test_dict = []
for i in range(len(test)):
test_dict.append( {'x': os.path.join(root, "Normal",test[i]), 'm': os.path.join(root, "Mask",test[i]), 'd': os.path.join(root, "Depth",test[i])})
return(train_dict, test_dict)
def make_filenames_json(train_dict, test_dict, root, postfix = ""):
with open(os.path.join(root, "train" + postfix+ ".json"), "w") as write_file:
json.dump(train_dict, write_file, indent=4)
with open(os.path.join(root, "test" + postfix+ ".json"), "w") as write_file:
json.dump(test_dict, write_file, indent=4)
if __name__ == "__main__":
im_path= '/network/tmp1/ccai/data/munit_dataset/simdata/Unity1000R_fL_lowRes/'
save_path = '/network/tmp1/ccai/data/mask_generation/Unity1000R_fL_lowRes/'#'/network/tmp1/ccai/data/mask_generation/WD/'
main(im_path, save_path, train_r = 0.8, postfix = "")
train_dict, test_dict = get_train_test_dict(im_path, 0.8)
#make_filenames_json(train_dict, test_dict, save_path)
make_filenames_json(train_dict, test_dict, save_path, postfix = "")