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