backup
This commit is contained in:
+91
-65
@@ -21,7 +21,7 @@ from data.render import Render, DepthRenderer
|
||||
from superSMPL.IO import readPC2Frame, writePC2
|
||||
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera
|
||||
from cloth3d.read import DataReader
|
||||
from cloth3d.util import quads2tris
|
||||
from cloth3d.util import quads2tris, zRotMatrix
|
||||
|
||||
|
||||
class Dataset(Dataset):
|
||||
@@ -55,21 +55,23 @@ class Dataset(Dataset):
|
||||
path = os.path.join(self.folder, frame.id)
|
||||
info = self.reader.read_info(frame.id)
|
||||
zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
|
||||
trans = info['trans'][:,frame.frame_nr]
|
||||
trans = info['trans'][:,frame.frame_nr].astype(np.float32)
|
||||
# mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
|
||||
|
||||
mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
|
||||
mesh_posed = torch.tensor(mesh_posed, device=self.device)
|
||||
# mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False)
|
||||
mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
|
||||
mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=self.device)
|
||||
|
||||
mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
|
||||
# TODO: what's the difference between masking and reading .obj file. is the obj file the not registered version?
|
||||
faces = mask_faces(self.body_faces, mesh_mask)
|
||||
# faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
|
||||
# body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
|
||||
body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, trans, zrot)
|
||||
# body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
|
||||
body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
|
||||
body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
|
||||
body = torch.tensor(body.astype(np.float32), device=self.device)
|
||||
body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=self.device)
|
||||
body_faces = quads2tris(body_faces)
|
||||
|
||||
if self.random_view:
|
||||
@@ -78,14 +80,30 @@ class Dataset(Dataset):
|
||||
z_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi)
|
||||
view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((x_angle, y_angle, z_angle)))
|
||||
|
||||
zRot = torch.tensor(zRotMatrix(-zrot), device=self.device)
|
||||
mesh_posed = torch.matmul(zRot, mesh_posed.T).T
|
||||
mesh_posed = mesh_posed - torch.tensor(trans, device=self.device)
|
||||
body = torch.matmul(zRot, body.T).T
|
||||
body = body - torch.tensor(trans, device=self.device)
|
||||
|
||||
transform = pytorch3d.transforms.Rotate(view)
|
||||
mesh_posed = transform.transform_points(mesh_posed)
|
||||
mesh_posed_no_trans = transform.transform_points(mesh_posed_no_trans)
|
||||
body = transform.transform_points(body)
|
||||
body_no_trans = transform.transform_points(body_no_trans)
|
||||
|
||||
|
||||
zRot = torch.tensor(zRotMatrix(zrot), device=self.device)
|
||||
mesh_posed = mesh_posed + torch.tensor(trans, device=self.device)
|
||||
mesh_posed = torch.matmul(zRot, mesh_posed.T).T
|
||||
body = body + torch.tensor(trans, device=self.device)
|
||||
body = torch.matmul(zRot, body.T).T
|
||||
|
||||
else:
|
||||
view = torch.eye(3, device=self.device)
|
||||
|
||||
combined = torch.concat((body, mesh_posed[mesh_mask, :]))
|
||||
# combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda')
|
||||
combined_no_trans = torch.concat((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))
|
||||
combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(self.device)
|
||||
|
||||
|
||||
@@ -99,11 +117,13 @@ class Dataset(Dataset):
|
||||
|
||||
start = time.perf_counter()
|
||||
depth_image, depth_mask, v2i, K, camera, pix_to_face = self.render(combined, combined_faces, info['camLoc'])
|
||||
# cam_no_trans = info['camLoc'].copy()
|
||||
# cam_no_trans[1:] = 0
|
||||
# depth_image_no_trans, depth_mask_no_trans, _, _, camera_no_trans, _ = self.render(combined_no_trans, combined_faces, cam_no_trans)
|
||||
cam_no_trans = info['camLoc'].copy()
|
||||
cam_no_trans[1:] = 0
|
||||
depth_image_no_trans, depth_mask_no_trans, _, _, camera_no_trans, _ = self.render(combined_no_trans, combined_faces, cam_no_trans)
|
||||
# camera_no_trans, _ = get_camera(cam_no_trans, [480, 640], self.device)
|
||||
print(f'pytorch3d time: {time.perf_counter() - start}')
|
||||
|
||||
|
||||
start = time.perf_counter()
|
||||
# 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)
|
||||
@@ -113,74 +133,78 @@ class Dataset(Dataset):
|
||||
depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask
|
||||
|
||||
|
||||
# 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()
|
||||
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)
|
||||
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 = 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())
|
||||
# right = torch.tensor([0], device=bbox.get_device())
|
||||
# top = torch.tensor([0], device=bbox.get_device())
|
||||
# bottom = torch.tensor([0], device=bbox.get_device())
|
||||
# if offset[0] > 0:
|
||||
# left = 2 * offset[0]
|
||||
# else:
|
||||
# right = -1 * 2 * offset[0]
|
||||
# if offset[1] > 0:
|
||||
# top = 2 * offset[1]
|
||||
# else:
|
||||
# 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 = 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()
|
||||
# new = centered[(centered.shape[0] // 2), (centered.shape[1] // 2)].cpu().numpy()
|
||||
# assert new == orig, 'Center pad did not work'
|
||||
left = torch.tensor([0], device=bbox.get_device())
|
||||
right = torch.tensor([0], device=bbox.get_device())
|
||||
top = torch.tensor([0], device=bbox.get_device())
|
||||
bottom = torch.tensor([0], device=bbox.get_device())
|
||||
if offset[0] > 0:
|
||||
left = 2 * offset[0]
|
||||
else:
|
||||
right = -1 * 2 * offset[0]
|
||||
if offset[1] > 0:
|
||||
top = 2 * offset[1]
|
||||
else:
|
||||
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 = 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()
|
||||
new = centered[(centered.shape[0] // 2), (centered.shape[1] // 2)].cpu().numpy()
|
||||
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,:,:]
|
||||
# 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)
|
||||
# upscaled = torch.nn.functional.pad(scaled, (left, right, top, bottom), value=-1)
|
||||
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)
|
||||
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)
|
||||
|
||||
|
||||
# Testing
|
||||
# Verifying
|
||||
|
||||
# pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True)
|
||||
# pc = pc[0,:,depth_mask].cpu().numpy()
|
||||
|
||||
# all_faces = combined_faces
|
||||
# visible_faces_idx = pix_to_face * depth_mask
|
||||
# visible_faces_idx = visible_faces_idx[torch.nonzero(visible_faces_idx, as_tuple=True)]
|
||||
# visible_vertices = torch.unique(torch.flatten(all_faces[visible_faces_idx, :])).to(torch.long)
|
||||
# visible_vertices_mask = (v2i[visible_vertices, 1] < depth_image.shape[0]) * (v2i[visible_vertices, 0] < depth_image.shape[1])
|
||||
# visible_vertices = visible_vertices[visible_vertices_mask]
|
||||
# visible_pixels_mask = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] != -1
|
||||
# visible_vertices = visible_vertices[visible_pixels_mask]
|
||||
|
||||
all_faces = combined_faces
|
||||
visible_faces_idx = pix_to_face * depth_mask
|
||||
visible_faces_idx = visible_faces_idx[torch.nonzero(visible_faces_idx, as_tuple=True)]
|
||||
visible_vertices = torch.unique(torch.flatten(all_faces[visible_faces_idx, :])).to(torch.long)
|
||||
visible_vertices_mask = (v2i[visible_vertices, 1] < depth_image.shape[0]) * (v2i[visible_vertices, 0] < depth_image.shape[1])
|
||||
visible_vertices = visible_vertices[visible_vertices_mask]
|
||||
visible_pixels_mask = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] != -1
|
||||
visible_vertices = visible_vertices[visible_pixels_mask]
|
||||
# x = combined[visible_vertices, :]
|
||||
# y = np.array(info['camLoc'], dtype=np.float32)
|
||||
# dist = torch.sqrt(torch.sum((x - torch.tensor(y[None,:], device=x.get_device())) ** 2, axis=-1))
|
||||
# dist_img = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]]
|
||||
# diff = torch.mean(torch.abs(dist - dist_img))
|
||||
|
||||
x = combined[visible_vertices, :]
|
||||
y = np.array(info['camLoc'], dtype=np.float32)
|
||||
dist = torch.sqrt(torch.sum((x - torch.tensor(y[None,:], device=x.get_device())) ** 2, axis=-1))
|
||||
dist_img = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]]
|
||||
diff = torch.mean(torch.abs(dist - dist_img))
|
||||
|
||||
# Plotting
|
||||
|
||||
# combined = combined.cpu().numpy()
|
||||
# depth_image = depth_image.cpu().numpy()
|
||||
@@ -219,6 +243,7 @@ class Dataset(Dataset):
|
||||
# image = self.transforms(image)
|
||||
data = { 'mesh_posed': mesh_posed,
|
||||
'depth_image': depth_image,
|
||||
# 'depth_image': upscaled,
|
||||
'mesh_mask': mesh_mask,
|
||||
'depth_mask': depth_mask,
|
||||
'normal_image': normal_image,
|
||||
@@ -226,6 +251,7 @@ class Dataset(Dataset):
|
||||
'depth_mask_body': depth_mask_body,
|
||||
'rgb_image': rgb_image.to(self.device),
|
||||
'view': view,
|
||||
'info': info,
|
||||
}
|
||||
|
||||
self.cache[frame] = data
|
||||
@@ -238,7 +264,7 @@ if __name__ == '__main__':
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# dataset = Dataset('/home/workspace/local-data/train')
|
||||
dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
|
||||
dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt', random_view=True)
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
|
||||
for batch in train_dataloader:
|
||||
|
||||
Reference in New Issue
Block a user