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}