def get_data(root_folder, filenamebase, color_on_depth_suffix='_color_on_depth.png', depth_suffix='_depth.png', img_res=(320, 240), as_torch=True): # load color color_on_depth_image_filepath = root_folder + filenamebase + color_on_depth_suffix color_on_depth_image = read_RGB_image(color_on_depth_image_filepath, new_res=img_res) # load depth filenamebase_split = filenamebase.split('/') depth_filenamebase = '/'.join( filenamebase_split[0:2]) + '/depth/' + filenamebase_split[-1] depth_image_filepath = root_folder + depth_filenamebase + depth_suffix depth_image = read_RGB_image(depth_image_filepath, new_res=img_res) depth_image = np.array(depth_image) depth_image = np.reshape(depth_image, (depth_image.shape[0], depth_image.shape[1], 1)) # get data RGBD_image = np.concatenate((color_on_depth_image, depth_image), axis=-1) RGBD_image = RGBD_image.swapaxes(1, 2).swapaxes(0, 1) img_data = RGBD_image if as_torch: img_data = torch.from_numpy(RGBD_image).float() return img_data
def __getitem__(self, idx): colour = io_image.read_RGB_image( self.colour_filepaths[self.file_idxs[idx]], new_res=self.img_res) depth = io_image.read_RGB_image( self.depth_filepaths[self.file_idxs[idx]], new_res=self.img_res) depth = np.reshape(depth, (depth.shape[0], depth.shape[1], 1)) RGBD_image = np.concatenate((colour, depth), axis=-1).astype(float) RGBD_image = np.divide(RGBD_image, 255) RGBD_image = RGBD_image.swapaxes(1, 2).swapaxes(0, 1) RGBD_image = torch.from_numpy(RGBD_image).float() return RGBD_image, 0
def __getitem__(self, idx): pose = genfromtxt(self.pose_filepaths[self.file_idxs[idx]], delimiter=',')[1:] pose = torch.from_numpy(pose).float() colour = io_image.read_RGB_image( self.colour_filepaths[self.file_idxs[idx]]) mask = io_image.read_RGB_image( self.mask_filepaths[self.file_idxs[idx]]) cropped_img = self.crop_image(colour, mask) colour = io_image.change_res_image(colour, self.img_res) mask = io_image.change_res_image(mask, self.img_res) cropped_img = io_image.change_res_image(cropped_img, self.img_res) with_imagenet = io_image.read_RGB_image( self.with_imagenet_filepaths[self.file_idxs[idx]], new_res=self.img_res) data_image = with_imagenet if (not self.noise_channel) and self.num_channels > 3: depth = io_image.read_RGB_image( self.depth_filepaths[self.file_idxs[idx]], new_res=self.img_res) depth = np.reshape(depth, (depth.shape[0], depth.shape[1], 1)) data_image = np.concatenate((data_image, depth), axis=-1).astype(float) #cropped_img_non_noisy = np.copy(cropped_img) #cropped_img_noisy, noise_idxs = util.add_noise(cropped_img, 0.2) data_image_noisy, noise_idxs = util.add_noise(data_image, self.noise_level) #colour = np.concatenate((colour, noise_idxs), axis=-1).astype(float) if self.transform: data_image_noisy = self.transform(data_image_noisy).float() noise_idxs = self.transform(noise_idxs).float() #cropped_img_noisy = self.transform(cropped_img_noisy) #cropped_img_non_noisy = self.transform(cropped_img_non_noisy) colour = self.transform(colour).float() data_image_noisy = torch.cat((data_image_noisy, noise_idxs), 0) #vis.plot_image((data_image_noisy.numpy()[0:3, :, :] + 0.5) * 255) #vis.show() #vis.plot_image((colour.numpy() + 0.5) * 255) #vis.show() return data_image_noisy, colour
def read_unreal_depth_img(self, img_filepath, unreal_max_depth=100): ''' Unreal max depth should be in cm ''' img = io_image.read_RGB_image(img_filepath) img = img[:, :, 0] + img[:, :, 1] / 255 + img[:, :, 2] / 65025 # assuming img[0, 0] contains the maximum available depth in the unreal image # as in Unreal I put a large cuboid at height 100 cm to calibrate the depth img *= unreal_max_depth / img[0, 0] # put 0 on every pixel that has max depth img[np.abs(img - unreal_max_depth) < 1e-2] = 0 # convert to mm img *= 10 return img
def read_image_from_image_grid(img_filepath, img_res, grid_x, grid_y): img_grid = io_image.read_RGB_image((img_filepath)) img = np.zeros(img_res) for i in range(3): start_y = grid_x * img_res[0] start_x = grid_y * img_res[1] img += img_grid[start_x:start_x + img_res[0], start_y:start_y + img_res[1], i] * np.power(255, i) img = img / 255 return img
def _get_data(root_folder, filenamebase, new_res, as_torch=True, depth_suffix='_depth.png', color_on_depth_suffix='_color_on_depth.png'): # load color color_on_depth_image_filename = root_folder + filenamebase + color_on_depth_suffix color_on_depth_image = read_RGB_image(color_on_depth_image_filename, new_res=new_res) # load depth depth_image_filename = root_folder + filenamebase + depth_suffix depth_image = read_RGB_image(depth_image_filename, new_res=new_res) depth_image = np.array(depth_image) depth_image = np.reshape(depth_image, (depth_image.shape[0], depth_image.shape[1], 1)) # get data RGBD_image = np.concatenate((color_on_depth_image, depth_image), axis=-1) RGBD_image = RGBD_image.swapaxes(1, 2).swapaxes(0, 1) data = RGBD_image if as_torch: data = torch.from_numpy(RGBD_image).float() return data
results_dir = './results/' k = 256 hidden = 128 num_channels_in = 4 num_channels_out = 3 img_path = 'C:/Users/Administrator/Documents/Datasets/ycb_unreal_colour (493).png' img_save_path = results_dir + 'output_img.png' img_res = (640, 480) noise_level = 0.01 transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize( (0.5, 0.5, 0.5, 0.5), (0.5, 0.5, 0.5, 0.5))]) # load image input_image = io_image.read_RGB_image(img_path, img_res) if input_image.shape[2] == 4: input_image = input_image[:, :, 0:3] data_image_noisy, noise_idxs = util.add_noise(input_image, noise_level) data_image_noisy = transform(data_image_noisy).float() noise_idxs = transform(noise_idxs).float() data_image_noisy = torch.cat((data_image_noisy, noise_idxs), 0) # load model checkpoint = torch.load(results_dir + 'ycb_checkpoint.pth.tar') args = checkpoint['args'] epoch = checkpoint['epoch'] model = VQ_CVAE(d=hidden, k=k, num_channels_in=num_channels_in, num_channels_out=num_channels_out) model.load_state_dict(checkpoint['state_dict']) if use_cuda:
def read_depth_img(filepath): depth_img = io_image.read_RGB_image(filepath) return depth_img
def read_color_img(filepath, img_res=None): color_img = io_image.read_RGB_image(filepath, new_res=img_res) return color_img