def augment_3(self, image, depth, lidar, crop_size, degree):
# rsz = iaa.Resize({"height": resize_size[0], "width": resize_size[1]})
seq = iaa.Sequential(
[
iaa.PadToFixedSize(height=crop_size[0],
width=crop_size[1]), # 保证可crop
iaa.CropToFixedSize(height=crop_size[0],
width=crop_size[1]), # random crop
iaa.Fliplr(0.5),
# iaa.Flipud(0.5),
iaa.Rotate((-degree, degree)),
iaa.GammaContrast((0.9, 1.1)),
iaa.Multiply((0.9, 1.1)),
],
random_order=True)
depth, lidar = np.expand_dims(depth, 2), np.expand_dims(lidar, 2)
tmp = np.concatenate((depth, lidar), axis=2)
tmp = (tmp * 1000).astype(np.int32) # 米单位*1000保留精度
tmp = SegmentationMapsOnImage(tmp, shape=tmp.shape)
# image, tmp = rsz(image=image, segmentation_maps=tmp)
image, tmp = seq(image=image, segmentation_maps=tmp)
tmp = tmp.arr
tmp = tmp.astype(np.float32) / 1000 # 再转回米
depth, lidar = tmp[:, :, 0], tmp[:, :, 1]
return image, depth, lidar
def __getitem__(self, idx):
image = Image.open(os.path.join(
self.root, self.image_filenames[idx])).convert('RGB')
image = np.asarray(image)
if self.nums is not None:
segmap = np.zeros((image.shape[0], image.shape[1], 1),
dtype=np.int32)
nums = self.nums[idx]
for num in nums:
poly = Polygon(num['box'])
segmap = poly.draw_on_image(segmap, color=2)
segmap = SegmentationMapsOnImage(segmap, shape=image.shape)
if self.transforms:
image, segmap = self.transforms(image=image,
segmentation_maps=segmap)
image = to_tensor(image)
image = normalize(image, MEAN, STD)
segmap = segmap.arr
segmap = to_tensor(segmap.astype(np.float32))
return {'image': image, 'mask': segmap}
else:
if self.transforms:
image = self.transforms(image=image)
image = to_tensor(image)
image = normalize(image, MEAN, STD)
return {'image': image}
