-
Notifications
You must be signed in to change notification settings - Fork 5
/
split_dataset.py
52 lines (38 loc) · 1.42 KB
/
split_dataset.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
#! encoding: UTF-8
import os
import glob
import json
import subprocess
genome_VG_100K_path = "./data/genome/VG_100K"
genome_VG_100K_2_path = "./data/genome/VG_100K_2"
train_json_path = "./data/genome/train_split.json"
val_json_path = "./data/genome/val_split.json"
test_json_path = "./data/genome/test_split.json"
train_savepath = "./data/genome/im2p_train/"
val_savepath = "./data/genome/im2p_val/"
test_savepath = "./data/genome/im2p_test/"
os.mkdir(train_savepath)
os.mkdir(val_savepath)
os.mkdir(test_savepath)
all_images = glob.glob(genome_VG_100K_path + "/*.jpg")
all_images_2 = glob.glob(genome_VG_100K_2_path + "/*.jpg")
for item in all_images_2:
all_images.append(item)
with open(train_json_path, "r") as fr1:
train_names = json.load(fr1)
with open(val_json_path, "r") as fr2:
val_names = json.load(fr2)
with open(test_json_path, "r") as fr3:
test_names = json.load(fr3)
print "train images num: {}".format(len(train_names))
print "val images num: {}".format(len(val_names))
print "test images num: {}".format(len(test_names))
for idx, img in enumerate(all_images):
#print "idx: {} {}".format(idx, img)
img_name = int(os.path.basename(img).split(".")[0])
if img_name in train_names:
subprocess.call(["cp", img, train_savepath])
elif img_name in val_names:
subprocess.call(["cp", img, val_savepath])
elif img_name in test_names:
subprocess.call(["cp", img, test_savepath])