-
Notifications
You must be signed in to change notification settings - Fork 8
/
data_loader_v2.py
143 lines (121 loc) · 4.46 KB
/
data_loader_v2.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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
import os
from PIL import Image
import numpy as np
import os.path as osp
import io
import random
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from torchvision.transforms import *
from torchvision.transforms import functional as F
def read_image(img_path):
"""Keep reading image until succeed.
This can avoid IOError incurred by heavy IO process."""
got_img = False
if not osp.exists(img_path):
raise IOError("{} does not exist".format(img_path))
while not got_img:
try:
img = Image.open(img_path).convert('RGB')
got_img = True
except IOError:
print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
pass
return img
class ImageDatasettest(Dataset):
"""Image Person ReID Dataset"""
def __init__(self, dataset, transform=None):
self.dataset = dataset
self.transform = transform
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
img_path, pid, camid = self.dataset[index]
img = read_image(img_path)
# print(img_path)
if self.transform is not None:
img = self.transform(img)
return img, pid, camid
class ImageDatasettrain(Dataset):
"""Image Person ReID Dataset"""
def __init__(self, dataset, height, width):
self.dataset = dataset
self.height=height
self.width=width
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
img_path, pid, _ = self.dataset[index]
img1 = read_image(img_path)
path2=img_path.split('/')[0].replace('mlr_duke','duke')
#path2=img_path.split('/')[0].replace('mlr_cuhk03','cuhk03')
#path2=img_path.split('/')[0].replace('mlr_market1501','market1501')
#path2=img_path.split('/')[0].replace('mlr_vp_v1','vp')
#path2=path2.replace('bmp','png')
#path2=img_path.split('/')[0].replace('mlr_caviar','caviar')
img2=read_image(path2)
width, height = img1.size
resolution=(width*1.0)/self.width
Random2DT=Random2DTranslation(self.height,self.width)
RandomHor=RandomHorizontalFlip2()
Pad = transforms.Pad(10)
Resize = transforms.Resize([256,128])
toten=ToTensor()
normm=Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
img1 = Resize(img1)
img2 = Resize(img2)
img1,img2=RandomHor(img1,img2)
img1 = Pad(img1)
img2 = Pad(img2)
img1,img2=Random2DT.twoimages(img1,img2)
img1,img2=toten(img1),toten(img2)
img1,img2=normm(img1),normm(img2)
return img1, img2, pid ,img_path ,path2
class RandomHorizontalFlip2(object):
def __init__(self, p=0.5):
self.p = p
def __call__(self, img1,img2):
if random.random() < self.p:
return F.hflip(img1),F.hflip(img2)
return img1,img2
class Random2DTranslation(object):
"""
With a probability, first increase image size to (1 + 1/8), and then perform random crop.
Args:
- height (int): target height.
- width (int): target width.
- p (float): probability of performing this transformation. Default: 0.5.
"""
def __init__(self, height, width, p=1, interpolation=Image.BILINEAR):
self.height = height
self.width = width
self.p = p
self.interpolation = interpolation
def __call__(self, img):
"""
Args:
- img (PIL Image): Image to be cropped.
"""
#resized_img = img.resize((new_width, new_height), self.interpolation)
x_maxrange = 20
y_maxrange = 20
x1 = int(round(random.uniform(0, x_maxrange)))
y1 = int(round(random.uniform(0, y_maxrange)))
croped_img = img.crop((x1, y1, x1 + self.width, y1 + self.height))
return croped_img
def twoimages(self, img1,img2):
"""
Args:
- img (PIL Image): Image to be cropped.
"""
x_maxrange = 20
y_maxrange = 20
x1 = int(round(random.uniform(0, x_maxrange)))
y1 = int(round(random.uniform(0, y_maxrange)))
croped_img1 = img1.crop((x1, y1, x1 + self.width, y1 + self.height))
croped_img2 = img2.crop((x1, y1, x1 + self.width, y1 + self.height))
return croped_img1, croped_img2