def __getitem__(self, index):
img_path = self.base_dir / self.img_paths[index]
bbox = change_bbox(self.bbox[index], 1.4, use_forehead=False)
img = np.array(Image.open(img_path).crop(bbox))
label = self.labels[index].copy()
if self.use_bined:
bined_label = self.euler_binned[index].copy()
# ImageAugment (RandomBrightness, AddNoise...)
if self.image_augmenter:
augmented = self.image_augmenter(image=img)
img = augmented['image']
# Resize (Scale & Pad & Crop)
if self.resizer:
resized = self.resizer(image=img)
img = resized['image']
# AffineAugment (Horizontal Flip, Rotate...)
if self.affine_augmenter:
augmented = self.affine_augmenter(image=img)
img = augmented['image']
if self.split=='train':
# 图片左右翻转
if random.random() < 0.5:
img = cv2.flip(img, 1)
label[0] = -label[0]
label[2] = -label[2]
if self.use_bined:
bined_label[0] = -(bined_label[0]-3)+3
bined_label[2] = -(bined_label[2]-9)+10
if random.random() < 0.5 and abs(label[0])<30 and abs(label[2])<30:
if random.random() < 0.5:
factor = 1
label[2] += 90
if self.use_bined:
bined_label[2] = min(bined_label[2] + 10, 20)
else:
factor = 3
label[2] -= 90
if self.use_bined:
bined_label[2] = max(bined_label[2] - 10, 0)
img = np.ascontiguousarray(np.rot90(img, factor))
if self.n_class == 4:
label = euler2quat(*label)
if self.debug:
print(self.img_paths[index], label)
else:
img = preprocess(img)
img = img.transpose(2, 0, 1)
img = torch.FloatTensor(img)
label = torch.FloatTensor(label)
if self.use_bined:
return img, label, bined_label[0], bined_label[1], bined_label[2]
else:
return img, label
def __getitem__(self, index):
index = index % len(self.ids)
idxs = np.random.choice(self.ids_index[index], size=2, replace=False)
img_path1 = self.base_dir / (self.ids[index] + '_%d.jpg' % idxs[0])
img_path2 = self.base_dir / (self.ids[index] + '_%d.jpg' % idxs[1])
# scale = np.random.random_sample() * 0.2 + 0.1
scale = np.random.random_sample() * 0.2 + 1.4
bbox1 = change_bbox(self.bboxs[index][idxs[0]],
scale=scale,
use_forehead=False)
bbox2 = change_bbox(self.bboxs[index][idxs[1]],
scale=scale,
use_forehead=False)
img1 = np.array(Image.open(img_path1).crop(bbox1))
img2 = np.array(Image.open(img_path2).crop(bbox2))
lbl1 = self.labels[index][idxs[0]]
lbl2 = self.labels[index][idxs[1]]
if self.use_bined:
bined_label = self.euler_binned[index].copy()
bined_lbl1 = bined_label[idxs[0]]
bined_lbl2 = bined_label[idxs[1]]
# ImageAugment (RandomBrightness, AddNoise...)
if self.image_augmenter:
augmented = self.image_augmenter(image=img1)
img1 = augmented['image']
augmented = self.image_augmenter(image=img2)
img2 = augmented['image']
# Resize (Scale & Pad & Crop)
if self.resizer:
resized = self.resizer(image=img1)
img1 = resized['image']
resized = self.resizer(image=img2)
img2 = resized['image']
# AffineAugment (Horizontal Flip, Rotate...)
if self.affine_augmenter:
augmented = self.affine_augmenter(image=img1)
img1 = augmented['image']
augmented = self.affine_augmenter(image=img2)
img2 = augmented['image']
# label = (lbl1 > lbl2) * 2 - 1
label = np.sign(lbl1 - lbl2)
if self.n_class == 4:
lbl1 = euler2quat(*lbl1)
lbl2 = euler2quat(*lbl2)
if self.debug:
print(label)
return img1, img2
else:
img1 = preprocess(img1)
img1 = torch.FloatTensor(img1).permute(2, 0, 1)
img2 = preprocess(img2)
img2 = torch.FloatTensor(img2).permute(2, 0, 1)
label = torch.FloatTensor(label.astype(np.float32))
lbl1 = torch.FloatTensor(lbl1)
lbl2 = torch.FloatTensor(lbl2)
if self.use_bined:
return img1, img2, lbl1, lbl2, label, int(bined_lbl1[0]), int(bined_lbl1[1]), int(bined_lbl1[2]),\
int(bined_lbl2[0]), int(bined_lbl2[1]), int(bined_lbl2[2])
else:
return img1, img2, lbl1, lbl2, label