def test_it_should_add_sample(self): dataset = Dataset("data") dataset.add_pose(self.dummy_rgb, self.dummy_depth, self.dummy_pose) self.assertEqual(dataset.size(), 1)
vpRender = ModelRenderer(model["model_path"], SHADER_PATH, real_dataset.camera, window, window_size) vpRender.load_ambiant_occlusion_map(model["ambiant_occlusion_model"]) OBJECT_WIDTH = int(model["object_width"]) metadata = {} metadata["translation_range"] = str(TRANSLATION_RANGE) metadata["rotation_range"] = str(ROTATION_RANGE) metadata["image_size"] = str(IMAGE_SIZE[0]) metadata["save_type"] = data["save_type"] metadata["object_width"] = {} for model in MODELS: metadata["object_width"][model["name"]] = str(model["object_width"]) metadata["min_radius"] = str(SPHERE_MIN_RADIUS) metadata["max_radius"] = str(SPHERE_MAX_RADIUS) for i in range(real_dataset.size()): frame, pose = real_dataset.data_pose[i] rgb_render, depth_render = vpRender.render(pose.transpose()) rgb, depth = frame.get_rgb_depth(real_dataset.path) masked_rgb, masked_depth = mask_real_image(rgb, depth, depth_render) for j in range(SAMPLE_QUANTITY): rotated_rgb, rotated_depth, rotated_pose = random_z_rotation( masked_rgb, masked_depth, pose, real_dataset.camera) random_transform = Transform.random( (-TRANSLATION_RANGE, TRANSLATION_RANGE), (-ROTATION_RANGE, ROTATION_RANGE)) inverted_random_transform = Transform.from_parameters( *(-random_transform.to_parameters()))
class DataAugmentation: def __init__(self): self.occluder = None self.background = None self.rgb_noise = None self.depth_noise = None self.blur_kernel = None self.jitter = None self.hue_noise = None def set_background(self, path): self.background = RGBDDataset(path) def set_occluder(self, path): self.occluder = Dataset(path) self.occluder.load() def set_rgb_noise(self, gaussian_std): self.rgb_noise = gaussian_std def set_depth_noise(self, gaussian_std): self.depth_noise = gaussian_std def set_hue_noise(self, offset): """ offset is the % of random hue offset distribution :param offset: :return: """ self.hue_noise = offset def set_blur(self, size): self.blur_kernel = size def set_jitter(self, max_x, max_y): self.jitter = (max_x, max_y) def augment(self, rgb, depth, prior, real=False): ret_rgb = rgb ret_depth = depth if real and self.occluder: if random.uniform(0, 1) < 0.75: rand_id = random.randint(0, self.occluder.size() - 1) occluder_rgb, occluder_depth, occ_pose = self.occluder.load_image(rand_id) if random.randint(0, 1): occluder_rgb, occluder_depth, _ = self.occluder.load_pair(rand_id, 0) occluder_depth = occluder_depth.astype(np.float32) # Z offset of occluder to be closer to the occluded object ( with random distance in front of the object) offset = -occ_pose.matrix[2, 3] + prior.matrix[2, 3] - random.uniform(0.07, 0.01) occluder_depth += offset occluder_rgb = self.add_hue_noise(occluder_rgb, 1) occluder_rgb = imresize(occluder_rgb, ret_depth.shape, interp='nearest') occluder_depth = imresize(occluder_depth, ret_depth.shape, interp='nearest', mode="F").astype(np.int16) ret_rgb, ret_depth = self.depth_blend(ret_rgb, ret_depth, occluder_rgb, occluder_depth) if real and self.hue_noise: if random.uniform(0, 1) > 0.05: ret_rgb = self.add_hue_noise(ret_rgb, self.hue_noise) if self.jitter: self.x_jitter = random.randint(-self.jitter[0], self.jitter[0]) self.y_jitter = random.randint(-self.jitter[1], self.jitter[1]) if self.x_jitter > 0: ret_rgb = np.pad(ret_rgb, ((self.x_jitter, 0), (0, 0), (0, 0)), mode='constant')[:-self.x_jitter, :, :] ret_depth = np.pad(ret_depth, ((self.x_jitter, 0), (0, 0)), mode='constant')[:-self.x_jitter, :] else: ret_rgb = np.pad(ret_rgb, ((0, abs(self.x_jitter)), (0, 0), (0, 0)), mode='constant')[ abs(self.x_jitter):, :, :] ret_depth = np.pad(ret_depth, ((0, abs(self.x_jitter)), (0, 0)), mode='constant')[abs(self.x_jitter):, :] if self.y_jitter > 0: ret_rgb = np.pad(ret_rgb, ((0, 0), (self.y_jitter, 0), (0, 0)), mode='constant')[:, :-self.y_jitter, :] ret_depth = np.pad(ret_depth, ((0, 0), (self.y_jitter, 0)), mode='constant')[:, :-self.y_jitter] else: ret_rgb = np.pad(ret_rgb, ((0, 0), (0, abs(self.y_jitter)), (0, 0)), mode='constant')[:, abs(self.y_jitter):, :] ret_depth = np.pad(ret_depth, ((0, 0), (0, abs(self.y_jitter))), mode='constant')[:, abs(self.y_jitter):] if real and self.background: color_background, depth_background = self.background.load_random_image(ret_rgb.shape[1]) depth_background = depth_background.astype(np.int32) ret_rgb, ret_depth = self.color_blend(ret_rgb, ret_depth, color_background, depth_background) if real and self.rgb_noise: if random.uniform(0, 1) > 0.05: noise = random.uniform(0, self.rgb_noise) ret_rgb = self.add_noise(ret_rgb, noise) if real and self.depth_noise: if random.uniform(0, 1) > 0.05: noise = random.uniform(0, self.depth_noise) ret_depth = self.add_noise(ret_depth, noise) if real and self.blur_kernel is not None: if random.uniform(0, 1) < 0.75: kernel_size = random.randint(3, self.blur_kernel) kernel = self.gkern(kernel_size) ret_rgb[0, :, :] = scipy.signal.convolve2d(ret_rgb[0, :, :], kernel, mode='same') ret_rgb[1, :, :] = scipy.signal.convolve2d(ret_rgb[1, :, :], kernel, mode='same') ret_rgb[2, :, :] = scipy.signal.convolve2d(ret_rgb[2, :, :], kernel, mode='same') if random.uniform(0, 1) < 0.75: kernel_size = random.randint(3, self.blur_kernel) kernel = self.gkern(kernel_size) ret_depth[:, :] = scipy.signal.convolve2d(ret_depth[:, :], kernel, mode='same') return ret_rgb.astype(np.uint8), ret_depth @staticmethod def add_noise(img, gaussian_std): type = img.dtype copy = img.astype(np.float) gaussian_noise = np.random.normal(0, gaussian_std, img.shape) copy = (gaussian_noise + copy) if type == np.uint8: copy[copy < 0] = 0 copy[copy > 255] = 255 return copy.astype(type) @staticmethod def add_hue_noise(rgb, hue_offset): hsv = rgb2hsv(rgb) hsv[:, :, 0] = (hsv[:, :, 0] + random.uniform(-hue_offset, hue_offset)) % 1 rgb = hsv2rgb(hsv) * 255 return rgb.astype(np.uint8) @staticmethod def color_blend(rgb1, depth1, rgb2, depth2): mask = np.all(rgb1 == 0, axis=2) mask = ndimage.binary_dilation(mask).astype(mask.dtype) depth1[mask] = 0 rgb1[mask, :] = 0 mask = mask.astype(np.uint8) new_depth = depth2 * mask + depth1 new_color = rgb2 * mask[:, :, np.newaxis] + rgb1 return new_color.astype(np.uint8), new_depth @staticmethod def depth_blend(rgb1, depth1, rgb2, depth2): new_depth2 = depth2.copy() new_depth1 = depth1.copy() rgb1_mask = np.all(rgb1 == 0, axis=2) rgb2_mask = np.all(rgb2 == 0, axis=2) rgb1_mask = ndimage.binary_dilation(rgb1_mask) new_depth2[rgb2_mask] = -100000 new_depth1[rgb1_mask] = -100000 mask = (new_depth1 < new_depth2) pos_mask = mask.astype(np.uint8) neg_mask = (mask == False).astype(np.uint8) masked_rgb_occluder = rgb1 * pos_mask[:, :, np.newaxis] masked_rgb_object = rgb2 * neg_mask[:, :, np.newaxis] masked_depth_occluder = depth1 * pos_mask masked_depth_object = depth2 * neg_mask blend_rgb = masked_rgb_occluder + masked_rgb_object blend_depth = masked_depth_occluder + masked_depth_object return blend_rgb, blend_depth @staticmethod def gkern(kernlen=21, nsig=3.5): """Returns a 2D Gaussian kernel array.""" interval = (2 * nsig + 1.) / (kernlen) x = np.linspace(-nsig - interval / 2., nsig + interval / 2., kernlen + 1) kern1d = np.diff(st.norm.cdf(x)) kernel_raw = np.sqrt(np.outer(kern1d, kern1d)) kernel = kernel_raw / kernel_raw.sum() return kernel
def test_it_should_have_size_0_at_init(self): dataset = Dataset("data") self.assertEqual(dataset.size(), 0)
import cv2 import os if __name__ == '__main__': folder = "/home/mathieu/Dataset/DeepTrack/dragon/" dataset_path = os.path.join(folder, "train_raw_real") new_dataset_path = os.path.join(folder, "train_raw_real_resized") if not os.path.exists(new_dataset_path): os.mkdir(new_dataset_path) dataset = Dataset(dataset_path) if not dataset.load(): print("[Error]: Train dataset empty") sys.exit(-1) new_dataset = Dataset(new_dataset_path) new_dataset.camera = dataset.camera.copy() new_dataset.camera.set_ratio(2) for i in range(dataset.size()): rgb, depth, pose = dataset.load_image(i) new_rgb = cv2.resize( rgb, (new_dataset.camera.width, new_dataset.camera.height)) new_depth = cv2.resize( depth, (new_dataset.camera.width, new_dataset.camera.height)) new_dataset.add_pose(new_rgb, new_depth, pose) if i % (1 * dataset.size() / 100) == 0: print("Progress : {}%".format(i * 100 / dataset.size())) new_dataset.set_save_type(dataset.metadata["save_type"]) new_dataset.dump_images_on_disk() new_dataset.save_json_files(dataset.metadata)
elif key == NUM_PAD_7_KEY: detection_offset.rotate(z=math.radians(1)) elif key == NUM_PAD_8_KEY: detection_offset.translate(z=-0.001) elif key == NUM_PAD_9_KEY: detection_offset.rotate(x=math.radians(1)) elif key == ARROW_UP_KEY: detection_offset.translate(y=-0.001) elif key == ARROW_DOWN_KEY: detection_offset.translate(y=0.001) elif key == ARROW_LEFT_KEY: detection_offset.rotate(y=math.radians(-1)) elif key == ARROW_RIGHT_KEY: detection_offset.rotate(y=math.radians(1)) print("Compute detections") for i in range(dataset.size()): frame, pose = dataset.data_pose[i] # if pose is identity, compute the detection if pose == Transform(): rgb, depth = dataset.data_pose[i][0].get_rgb_depth(dataset.path) pose = detector.detect(rgb) if detector.get_likelihood() < 0.1: print( "[WARNING] : Detector returns uncertain pose at frame {}". format(i)) #Todo : need better way to handle viewpoint's pose change in dataset... dataset.data_pose[i] = (Frame(rgb, depth, str(i)), pose) np.save(os.path.join(dataset.path, "offset"), detection_offset.matrix) dataset.dump_images_on_disk() dataset.save_json_files({"save_type": "png"}) sensor.stop()
metadata["object_width"] = {} for model in MODELS: metadata["object_width"][model["name"]] = str(model["object_width"]) metadata["min_radius"] = str(SPHERE_MIN_RADIUS) metadata["max_radius"] = str(SPHERE_MAX_RADIUS) camera = Camera.load_from_json(data["camera_path"]) dataset = Dataset(OUTPUT_PATH, frame_class=data["save_type"]) dataset.camera = camera window_size = (camera.width, camera.height) window = InitOpenGL(*window_size) sphere_sampler = UniformSphereSampler(SPHERE_MIN_RADIUS, SPHERE_MAX_RADIUS) preload_count = 0 if PRELOAD: if dataset.load(): preload_count = dataset.size() print("This Dataset already contains {} samples".format( preload_count)) # Iterate over all models from config files for model in MODELS: vpRender = ModelRenderer(model["model_path"], SHADER_PATH, dataset.camera, window, window_size) vpRender.load_ambiant_occlusion_map(model["ambiant_occlusion_model"]) OBJECT_WIDTH = int(model["object_width"]) for i in tqdm(range(SAMPLE_QUANTITY - preload_count)): random_pose = sphere_sampler.get_random() random_transform = Transform.random( (-TRANSLATION_RANGE, TRANSLATION_RANGE), (-ROTATION_RANGE, ROTATION_RANGE)) pair = combine_view_transform(random_pose, random_transform)
occluder_path = "/home/mathieu/Dataset/DeepTrack/mixed/test" background_path = "/home/mathieu/Dataset/RGBD/SUN3D" object_dataset = Dataset(object_path) object_dataset.load() data_augmentation = DataAugmentation() data_augmentation.set_rgb_noise(2) data_augmentation.set_depth_noise(2) data_augmentation.set_hue_noise(0.07) data_augmentation.set_occluder(occluder_path) data_augmentation.set_background(background_path) data_augmentation.set_blur(5) # data_augmentation.set_jitter(20, 20) for i in range(object_dataset.size()): rgb, depth, pose = object_dataset.load_image(i) rgb, depth, label = object_dataset.load_pair(i, 0) rgb_augmented, depth_augmented = data_augmentation.augment( rgb, depth, pose, True) plt.figure(0) plt.imshow(rgb) plt.figure(1) plt.imshow(rgb_augmented) plt.figure(2) plt.imshow(rgb - rgb_augmented) plt.show() plt.figure(0) plt.imshow(depth)
model = MODELS[0] vpRender = ModelRenderer(model["model_path"], SHADER_PATH, real_dataset.camera, window) vpRender.load_ambiant_occlusion_map(model["ambiant_occlusion_model"]) OBJECT_WIDTH = int(model["object_width"]) metadata = {} metadata["translation_range"] = str(TRANSLATION_RANGE) metadata["rotation_range"] = str(ROTATION_RANGE) metadata["image_size"] = str(IMAGE_SIZE[0]) metadata["save_type"] = data["save_type"] metadata["object_width"] = {} for model in MODELS: metadata["object_width"][model["name"]] = str(model["object_width"]) metadata["min_radius"] = str(SPHERE_MIN_RADIUS) metadata["max_radius"] = str(SPHERE_MAX_RADIUS) for i in range(real_dataset.size()): frame, pose = real_dataset.data_pose[i] rgb_render, depth_render = vpRender.render(pose.transpose()) rgb, depth = frame.get_rgb_depth(real_dataset.path) masked_rgb, masked_depth = mask_real_image(rgb, depth, depth_render) for j in range(SAMPLE_QUANTITY): rotated_rgb, rotated_depth, rotated_pose = random_z_rotation(masked_rgb, masked_depth, pose, real_dataset.camera) random_transform = Transform.random((-TRANSLATION_RANGE, TRANSLATION_RANGE), (-ROTATION_RANGE, ROTATION_RANGE)) inverted_random_transform = Transform.from_parameters(*(-random_transform.to_parameters())) previous_pose = rotated_pose.copy() previous_pose = combine_view_transform(previous_pose, inverted_random_transform)