-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathpre_process_sysu.py
118 lines (95 loc) · 3.72 KB
/
pre_process_sysu.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
110
111
112
113
114
115
116
117
118
import numpy as np
from PIL import Image
import pdb
import os
# Modified https://github.com/mangye16/Cross-Modal-Re-ID-baseline/blob/master/pre_process_sysu.py
data_path = './SYSU-MM01/'
rgb_cameras = ['cam1','cam2','cam4','cam5']
ir_cameras = ['cam3','cam6']
# load id info
file_path_train = os.path.join(data_path,'exp/train_id.txt')
file_path_val = os.path.join(data_path,'exp/val_id.txt')
file_path_test = os.path.join(data_path,'exp/test_id.txt')
with open(file_path_train, 'r') as file:
ids = file.read().splitlines()
ids = [int(y) for y in ids[0].split(',')]
id_train = ["%04d" % x for x in ids]
with open(file_path_val, 'r') as file:
ids = file.read().splitlines()
ids = [int(y) for y in ids[0].split(',')]
id_val = ["%04d" % x for x in ids]
with open(file_path_test, 'r') as file:
ids = file.read().splitlines()
ids = [int(y) for y in ids[0].split(',')]
id_test = ["%04d" % x for x in ids]
# All VI-ReID baselines combine train and val split
id_train.extend(id_val)
# In case, you want to use train+val+test for source domain in Stage-2
# id_train.extend(id_test)
files_rgb = []
files_ir = []
print("Begin\n\n")
for id in sorted(id_train):
for cam in rgb_cameras:
img_dir = os.path.join(data_path,cam,id)
if os.path.isdir(img_dir):
new_files = sorted([img_dir+'/'+i for i in os.listdir(img_dir)])
files_rgb.extend(new_files)
for cam in ir_cameras:
img_dir = os.path.join(data_path,cam,id)
if os.path.isdir(img_dir):
new_files = sorted([img_dir+'/'+i for i in os.listdir(img_dir)])
files_ir.extend(new_files)
# test_files_rgb = []
# test_files_ir = []
# for id in sorted(id_test):
# for cam in rgb_cameras:
# img_dir = os.path.join(data_path,cam,id)
# if os.path.isdir(img_dir):
# new_files = sorted([img_dir+'/'+i for i in os.listdir(img_dir)])
# test_files_rgb.extend(new_files)
# for cam in ir_cameras:
# img_dir = os.path.join(data_path,cam,id)
# if os.path.isdir(img_dir):
# new_files = sorted([img_dir+'/'+i for i in os.listdir(img_dir)])
# test_files_ir.extend(new_files)
# relabel
pid_container = set()
for img_path in files_ir:
pid = int(img_path[-13:-9])
pid_container.add(pid)
print(len(pid_container))
pid2label = {pid:label for label, pid in enumerate(pid_container)}
fix_image_width = 144
fix_image_height = 288
def read_imgs(train_image):
train_img = []
train_label = []
for img_path in train_image:
# img
img = Image.open(img_path)
img = img.resize((fix_image_width, fix_image_height), Image.ANTIALIAS)
pix_array = np.array(img)
train_img.append(pix_array)
# label
pid = int(img_path[-13:-9])
pid = pid2label[pid]
train_label.append(pid)
return np.array(train_img), np.array(train_label)
# rgb imges
train_img, train_label = read_imgs(files_rgb)
np.save(data_path + 'train_rgb_resized_img.npy', train_img)
np.save(data_path + 'train_rgb_resized_label.npy', train_label)
# ir imges
train_img, train_label = read_imgs(files_ir)
np.save(data_path + 'train_ir_resized_img.npy', train_img)
np.save(data_path + 'train_ir_resized_label.npy', train_label)
# # rgb imges
# train_img, train_label = read_imgs(files_rgb)
# np.save(data_path + 'train_large_rgb_resized_img.npy', train_img)
# np.save(data_path + 'train_large_rgb_resized_label.npy', train_label)
# # ir imges
# train_img, train_label = read_imgs(files_ir)
# np.save(data_path + 'train_large_ir_resized_img.npy', train_img)
# np.save(data_path + 'train_large_ir_resized_label.npy', train_label)
print("Done creating .npy files")