-
Notifications
You must be signed in to change notification settings - Fork 43
/
CubBird_prepare_csv.py
109 lines (77 loc) · 3.34 KB
/
CubBird_prepare_csv.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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Wenbin Li
## Date: Dec. 16 2018
##
## Divide data into train/val/test in a csv version
## Output: train.csv, val.csv, test.csv
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
import os
import csv
import numpy as np
import random
from PIL import Image
import pdb
data_dir = '/FewShot/Datasets/CUB_birds' # the path of the download dataset
save_dir = '/FewShot/Datasets/CUB_birds/For_FewShot' # the saving path of the divided dataset
if not os.path.exists(os.path.join(save_dir, 'images')):
os.makedirs(os.path.join(save_dir, 'images'))
images_dir = os.path.join(data_dir, 'images')
train_class_num = 130
val_class_num = 20
test_class_num = 50
# get all the dog classes
classes_list = [class_name for class_name in os.listdir(images_dir) if os.path.isdir(os.path.join(images_dir, class_name))]
# divide the train/val/test set
random.seed(200)
train_list = random.sample(classes_list, train_class_num)
remain_list = [rem for rem in classes_list if rem not in train_list]
val_list = random.sample(remain_list, val_class_num)
test_list = [rem for rem in remain_list if rem not in val_list]
# save data into csv file----- Train
train_data = []
for class_name in train_list:
images = [[i, class_name] for i in os.listdir(os.path.join(images_dir, class_name))]
train_data.extend(images)
print('Train----%s' %class_name)
# read images and store these images
img_paths = [os.path.join(images_dir, class_name, i) for i in os.listdir(os.path.join(images_dir, class_name))]
for index, img_path in enumerate(img_paths):
img = Image.open(img_path)
img = img.convert('RGB')
img.save(os.path.join(save_dir, 'images', images[index][0]), quality=100)
with open(os.path.join(save_dir, 'train.csv'), 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['filename', 'label'])
writer.writerows(train_data)
# save data into csv file----- Val
val_data = []
for class_name in val_list:
images = [[i, class_name] for i in os.listdir(os.path.join(images_dir, class_name))]
val_data.extend(images)
print('Val----%s' %class_name)
# read images and store these images
img_paths = [os.path.join(images_dir, class_name, i) for i in os.listdir(os.path.join(images_dir, class_name))]
for index, img_path in enumerate(img_paths):
img = Image.open(img_path)
img = img.convert('RGB')
img.save(os.path.join(save_dir, 'images', images[index][0]), quality=100)
with open(os.path.join(save_dir, 'val.csv'), 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['filename', 'label'])
writer.writerows(val_data)
# save data into csv file----- Test
test_data = []
for class_name in test_list:
images = [[i, class_name] for i in os.listdir(os.path.join(images_dir, class_name))]
test_data.extend(images)
print('Test----%s' %class_name)
# read images and store these images
img_paths = [os.path.join(images_dir, class_name, i) for i in os.listdir(os.path.join(images_dir, class_name))]
for index, img_path in enumerate(img_paths):
img = Image.open(img_path)
img = img.convert('RGB')
img.save(os.path.join(save_dir, 'images', images[index][0]), quality=100)
with open(os.path.join(save_dir, 'test.csv'), 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['filename', 'label'])
writer.writerows(test_data)