luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

add new cropping/scaling

Browse Source
This commit is contained in:
luni
2023-01-18 23:18:41 +00:00
parent fd185d1b48
commit 9e1c827526
1 changed files with 72 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+72 -52
View File
@@ -88,9 +88,9 @@ class Dataset(Dataset):
transform = pytorch3d.transforms.Rotate(view) transform = pytorch3d.transforms.Rotate(view)
mesh_posed = transform.transform_points(mesh_posed) mesh_posed = transform.transform_points(mesh_posed)
mesh_posed_no_trans = transform.transform_points(mesh_posed_no_trans) # mesh_posed_no_trans = transform.transform_points(mesh_posed_no_trans)
body = transform.transform_points(body) body = transform.transform_points(body)
body_no_trans = transform.transform_points(body_no_trans) # body_no_trans = transform.transform_points(body_no_trans)
zRot = torch.tensor(zRotMatrix(zrot), device=self.device) zRot = torch.tensor(zRotMatrix(zrot), device=self.device)
@@ -127,60 +127,79 @@ class Dataset(Dataset):
start = time.perf_counter() start = time.perf_counter()
# TODO: is it sure I need this K, and not some transformed one? # TODO: is it sure I need this K, and not some transformed one?
normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K) normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K)
normal_image = torch.permute(normal_image[0], [2,0,1])
print(f'normals time: {time.perf_counter() - start}') print(f'normals time: {time.perf_counter() - start}')
depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask
top_left = root_joint.astype(np.float32)
root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2].round() top_left[:,1] -= 1.25; top_left[:,2] += 1.25
joints_pixel = camera.transform_points_screen(torch.tensor(joints.astype(np.float32), device=combined.get_device()))[0,:,:2] bottom_right = root_joint.astype(np.float32)
joints_no_trans_pixel = camera_no_trans.transform_points_screen(torch.tensor(joints_no_trans.astype(np.float32), device=combined.get_device()))[0,:,:2] bottom_right[:,1] += 1.25; bottom_right[:,2] -= 1.25
ratio = ((joints_pixel[1,:] - joints_pixel[2,:]) ** 2).sum().sqrt() / ((joints_no_trans_pixel[1,:] - joints_no_trans_pixel[2,:]) ** 2).sum().sqrt() corners = torch.tensor(np.concatenate((top_left, bottom_right), axis=0), device=self.device)
corners_pixel = camera.transform_points_screen(corners)[:,:2].round().to(torch.int64)
center_offset = torch.tensor([240, 320], device=root_joint_pixel.get_device()) - root_joint_pixel
center_offset_ratio = center_offset / torch.tensor([480, 640], device=root_joint_pixel.get_device()) cropper = lambda x: tv.transforms.functional.crop(x, corners_pixel[0,1], corners_pixel[0,0], corners_pixel[1,1] - corners_pixel[0,1], corners_pixel[1,0] - corners_pixel[0,0])
bbox = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:] resizer = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
width = bbox[2] - bbox[0]
height = bbox[3] - bbox[1]
width = width + (width % 2)
height = height + (height % 2)
dims = torch.tensor([width, height], device=width.get_device())
# crop_size = torch.max(dims)
# top = (root_joint_pixel[1] - height / 2).round().to(torch.int64)
# left = (root_joint_pixel[0] - width / 2).round().to(torch.int64)
cropped = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), height.to(torch.int64), width.to(torch.int64))
new_root_joint_pixel = root_joint_pixel - bbox[:2]
offset = dims / 2 - new_root_joint_pixel
bbox[:2] = bbox[:2] - offset
dims = dims + torch.abs(offset)
# cropped_2 = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), dims[1].to(torch.int64), dims[0].to(torch.int64)) depth_mask_garment = resizer(cropper(depth_mask_garment)[None,None,:,:])[0,0,:,:]
depth_mask = resizer(cropper(depth_mask)[None,None,:,:])[0,0,:,:]
scaled = resizer(cropper(depth_image)[None,None,:,:])[0,0,:,:] * depth_mask
normal_image = resizer(cropper(normal_image)[None,:,:,:])[0,:,:,:] * depth_mask[None,:,:]
root_dist = np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2))
scaled = root_dist - scaled
# root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2].round()
# joints_pixel = camera.transform_points_screen(torch.tensor(joints.astype(np.float32), device=combined.get_device()))[0,:,:2]
# joints_no_trans_pixel = camera_no_trans.transform_points_screen(torch.tensor(joints_no_trans.astype(np.float32), device=combined.get_device()))[0,:,:2]
# ratio = ((joints_pixel[1,:] - joints_pixel[2,:]) ** 2).sum().sqrt() / ((joints_no_trans_pixel[1,:] - joints_no_trans_pixel[2,:]) ** 2).sum().sqrt()
# center_offset = torch.tensor([240, 320], device=root_joint_pixel.get_device()) - root_joint_pixel
# center_offset_ratio = center_offset / torch.tensor([480, 640], device=root_joint_pixel.get_device())
# bbox = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
# width = bbox[2] - bbox[0]
# height = bbox[3] - bbox[1]
# width = width + (width % 2)
# height = height + (height % 2)
# dims = torch.tensor([width, height], device=width.get_device())
# # crop_size = torch.max(dims)
# # top = (root_joint_pixel[1] - height / 2).round().to(torch.int64)
# # left = (root_joint_pixel[0] - width / 2).round().to(torch.int64)
# cropped = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), height.to(torch.int64), width.to(torch.int64))
# new_root_joint_pixel = root_joint_pixel - bbox[:2]
# offset = dims / 2 - new_root_joint_pixel
# bbox[:2] = bbox[:2] - offset
# dims = dims + torch.abs(offset)
# # cropped_2 = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), dims[1].to(torch.int64), dims[0].to(torch.int64))
left = torch.tensor([0], device=bbox.get_device()) # left = torch.tensor([0], device=bbox.get_device())
right = torch.tensor([0], device=bbox.get_device()) # right = torch.tensor([0], device=bbox.get_device())
top = torch.tensor([0], device=bbox.get_device()) # top = torch.tensor([0], device=bbox.get_device())
bottom = torch.tensor([0], device=bbox.get_device()) # bottom = torch.tensor([0], device=bbox.get_device())
if offset[0] > 0: # if offset[0] > 0:
left = 2 * offset[0] # left = 2 * offset[0]
else: # else:
right = -1 * 2 * offset[0] # right = -1 * 2 * offset[0]
if offset[1] > 0: # if offset[1] > 0:
top = 2 * offset[1] # top = 2 * offset[1]
else: # else:
bottom = -1 * 2 * offset[1] # bottom = -1 * 2 * offset[1]
# centered = tv.transforms.functional.crop(cropped, new_top.to(torch.int64), new_left.to(torch.int64), new_height.to(torch.int64), new_width.to(torch.int64)) # # centered = tv.transforms.functional.crop(cropped, new_top.to(torch.int64), new_left.to(torch.int64), new_height.to(torch.int64), new_width.to(torch.int64))
centered = torch.nn.functional.pad(cropped, (left.to(torch.int64), right.to(torch.int64), top.to(torch.int64), bottom.to(torch.int64)), value=-1) # centered = torch.nn.functional.pad(cropped, (left.to(torch.int64), right.to(torch.int64), top.to(torch.int64), bottom.to(torch.int64)), value=-1)
orig = depth_image[root_joint_pixel.to(torch.int64)[1], root_joint_pixel.to(torch.int64)[0]].cpu().numpy() # orig = depth_image[root_joint_pixel.to(torch.int64)[1], root_joint_pixel.to(torch.int64)[0]].cpu().numpy()
new = centered[(centered.shape[0] // 2), (centered.shape[1] // 2)].cpu().numpy() # new = centered[(centered.shape[0] // 2), (centered.shape[1] // 2)].cpu().numpy()
assert new == orig, 'Center pad did not work' # assert new == orig, 'Center pad did not work'
scaled = tv.transforms.Resize((centered.shape[1] / ratio).round().int().item(), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)(centered[None,None,:,:])[0,0,:,:] # scaled = tv.transforms.Resize((centered.shape[1] / ratio).round().int().item(), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)(centered[None,None,:,:])[0,0,:,:]
left, right = math.floor((depth_image.shape[1] - scaled.shape[1]) / 2), round((depth_image.shape[1] - scaled.shape[1]) / 2) # left, right = math.floor((depth_image.shape[1] - scaled.shape[1]) / 2), round((depth_image.shape[1] - scaled.shape[1]) / 2)
top, bottom = math.floor((depth_image.shape[0] - scaled.shape[0]) / 2), round((depth_image.shape[0] - scaled.shape[0]) / 2) # top, bottom = math.floor((depth_image.shape[0] - scaled.shape[0]) / 2), round((depth_image.shape[0] - scaled.shape[0]) / 2)
upscaled = torch.nn.functional.pad(scaled, (left, right, top, bottom), value=-1) # upscaled = torch.nn.functional.pad(scaled, (left, right, top, bottom), value=-1)
# Verifying # Verifying
@@ -235,21 +254,22 @@ class Dataset(Dataset):
# v2i_coverage[v2i[:,0], v2i[:, 1]] = 1 # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
# ex.imshow(v2i_coverage).show() # ex.imshow(v2i_coverage).show()
rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32) rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32).to(self.device)
rgb_image = resizer(cropper(rgb_image)[None,:,:,:])[0,:,:,:] * depth_mask[None,:,:]
# view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), trans.astype(np.float32)), axis=0), device=self.device) # view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), trans.astype(np.float32)), axis=0), device=self.device)
# image = self.transforms(image) # image = self.transforms(image)
data = { 'mesh_posed': mesh_posed, data = { 'mesh_posed': mesh_posed,
'depth_image': depth_image, 'depth_image': scaled,
# 'depth_image': upscaled, # 'depth_image': upscaled,
'mesh_mask': mesh_mask, 'mesh_mask': mesh_mask,
'depth_mask': depth_mask, 'depth_mask': depth_mask,
'normal_image': normal_image, 'normal_image': normal_image,
'depth_mask_garment': depth_mask_garment, 'depth_mask_garment': depth_mask_garment,
'depth_mask_body': depth_mask_body, # 'depth_mask_body': depth_mask_body,
'rgb_image': rgb_image.to(self.device), 'rgb_image': rgb_image,
'view': view, 'view': view,
'info': info, 'info': info,
} }
@@ -264,8 +284,8 @@ if __name__ == '__main__':
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
# dataset = Dataset('/home/workspace/local-data/train') # dataset = Dataset('/home/workspace/local-data/train')
dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt', random_view=True) dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=False)
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) train_dataloader = DataLoader(dataset, batch_size=1, drop_last=True, shuffle=False)
for batch in train_dataloader: for batch in train_dataloader:
batch batch