diff --git a/data/dataset.py b/data/dataset.py index 1e278d3..cb4ef54 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -19,7 +19,7 @@ from torch.utils.data import Dataset sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') 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 utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info from cloth3d.read import DataReader from cloth3d.util import quads2tris, zRotMatrix @@ -54,6 +54,8 @@ class Dataset(Dataset): path = os.path.join(self.folder, frame.id) info = self.reader.read_info(frame.id) + info['zrot'] = torch.pi / 2 + torch.pi/4 + zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0] trans = info['trans'][:,frame.frame_nr].astype(np.float32) # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True) @@ -103,7 +105,7 @@ class Dataset(Dataset): view = torch.eye(3, device=self.device) combined = torch.concat((body, mesh_posed[mesh_mask, :])) - combined_no_trans = torch.concat((body_no_trans, mesh_posed_no_trans[mesh_mask, :])) + # 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) @@ -111,7 +113,7 @@ class Dataset(Dataset): start = time.perf_counter() # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0] # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0] - deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32)) + # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32)) # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined, faces=combined_faces.cpu().numpy().astype(np.int32)) print(f'deodr time: {time.perf_counter() - start}') @@ -119,7 +121,7 @@ class Dataset(Dataset): 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) + # 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}') @@ -139,18 +141,29 @@ class Dataset(Dataset): bottom_right[:,1] += 1.25; bottom_right[:,2] -= 1.25 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) + root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=self.device))[0,:2].round() 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]) resizer = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST) - - depth_mask_garment = resizer(cropper(depth_mask_garment)[None,None,:,:])[0,0,:,:] + # bbox_before = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:] + + depth_mask_garment = cropper(depth_mask_garment) + scale = depth_mask_garment.shape[0] / 256 + depth_mask_garment = resizer(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 + # bbox_after = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:] + # scale = (bbox_before[2] - bbox_before[0]) / (bbox_after[2] - bbox_after[0]) + + # root_dist = torch.tensor(np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2)), device=self.device) + root_dist = torch.tensor(np.abs(root_joint[0][0] - info['camLoc'][0]), device=self.device) + # scaled = root_dist - scaled + scaled = scaled - root_dist # 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] @@ -259,19 +272,24 @@ class Dataset(Dataset): # 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) + info = filter_info(info, frame.frame_nr) # image = self.transforms(image) data = { 'mesh_posed': mesh_posed, 'depth_image': scaled, # 'depth_image': upscaled, 'mesh_mask': mesh_mask, - 'depth_mask': depth_mask, + 'depth_mask': depth_mask, #depth_mask, depth_mask_garment + 'depth_mask_full': depth_mask, #depth_mask, 'normal_image': normal_image, - 'depth_mask_garment': depth_mask_garment, + # 'depth_mask_garment': depth_mask_garment, # 'depth_mask_body': depth_mask_body, 'rgb_image': rgb_image, 'view': view, 'info': info, + 'root_dist': root_dist, + 'scale': scale, + 'root_joint_pixel': root_joint_pixel, } self.cache[frame] = data diff --git a/speed.py b/speed.py new file mode 100644 index 0000000..132632a --- /dev/null +++ b/speed.py @@ -0,0 +1,13 @@ +import numpy as np +import time + +if __name__ == '__main__': + start = time.perf_counter() + a, b = np.random.rand(1024, 1024), np.random.rand(1024, 1024) + print(f'init time: {time.perf_counter() - start}') + + start = time.perf_counter() + + c = np.dot(a, b) + + print(f'dot time: {time.perf_counter() - start}') diff --git a/test.py b/test.py index 821eabc..241bbfc 100644 --- a/test.py +++ b/test.py @@ -25,14 +25,67 @@ def plot3d(depth, mask): import kornia import plotly.graph_objects as go from data.render import intrinsic - K = torch.tensor(intrinsic()[None, :, :], device='cuda') + K = torch.tensor(intrinsic(RES_X=256, RES_Y=256)[None, :, :], device='cuda') pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True) pc = pc[0,:,mask[0]].detach().cpu().numpy() + pc[2,:] -= 3.09 go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), ]).show() +def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel, device='cuda'): + import kornia + import plotly.graph_objects as go + import plotly.express as ex + from cloth3d.util import zRotMatrix + import pytorch3d + import numpy as np + from utils.data_utils import intrinsic, extrinsic, get_camera + import torchvision as tv + cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device()) + # K = torch.tensor(intrinsic()[None, :, :], device=device) + + def restore(img, scale, center, root_dist=0.0): + import numpy as np + resizer = tv.transforms.Resize((torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)), interpolation=tv.transforms.functional.InterpolationMode.NEAREST) + cropper = lambda img, off: tv.transforms.functional.crop(img, off[1], off[0], 480, 640) + + out = resizer(img + root_dist) + offset = (np.array(np.round([out.shape[2]/2, out.shape[2]/2]))[None,:] - center.cpu().numpy()).astype(np.int64) + out = torch.stack([cropper(out[i], offset[i]) for i in range(img.shape[0])]) + out = out + + return out + + pc = cameras.unproject_points(depth) + mask = restore(mask[:,None,:,:], scale, root_joint_pixel).to(torch.bool) + depth = restore(depth, scale, root_joint_pixel, root_dist) * mask + pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True) + pc = pc[0,:,mask[0,0,:,:]] + pc[2,:] -= root_dist + pc = pc + pc = pc.detach().cpu().numpy() + + mesh = mesh[0] + # zRot = torch.tensor(zRotMatrix(-info['zrot']), device=device) + # mesh = torch.matmul(zRot, mesh.T).T + # mesh = mesh - info['trans'][0].to(torch.float32).to(device) + # zRot = torch.tensor(zRotMatrix(info['zrot']), device=device) + # mesh = torch.matmul(zRot, mesh.T).T + + # view = pytorch3d.transforms.axis_angle_to_matrix(torch.tensor([0.0, 0.0, torch.pi/2], device=device)) + # transform = pytorch3d.transforms.Rotate(view) + # mesh = transform.transform_points(mesh) + + mesh = mesh.detach().cpu().numpy() + + go.Figure(data = [go.Scatter3d(x=pc[2,:], y=pc[0,:], z=-pc[1,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), + go.Scatter3d(x=-mesh[:,0], y=mesh[:,1], z=mesh[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'), + ]).show() + + pass + # temp def plot_gt(depth, gt, mask): import kornia @@ -126,9 +179,10 @@ def train_step(batch): chamfer_reconst = depth_reconst_loss(y_pred, batch['depth_image'][:,None,:,:], batch['depth_mask']) normals_reconst = normals_reconst_loss(y_pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask']) # loss = mse #+ 1e-1 * (chamfer + normal) - loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst + loss = 1 - msssim + mse + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst # plot_gt(y_pred, batch) + plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel']) # 3. Zero grad optimizer optimizer.zero_grad() @@ -157,7 +211,7 @@ def valid_step(batch): chamfer_reconst = depth_reconst_loss(pred, batch['depth_image'][:,None,:,:], batch['depth_mask']) normals_reconst = normals_reconst_loss(pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask']) # loss = mse #+ 1e-1 * (chamfer + normal) - loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + normals_reconst + loss = 1 - msssim + mse + mse_normals #+ chamfer_reconst + mse + normals_reconst return { 'valid/loss': loss.detach().cpu().numpy() } diff --git a/utils/data_utils.py b/utils/data_utils.py index e9ecbaa..8644e35 100644 --- a/utils/data_utils.py +++ b/utils/data_utils.py @@ -63,9 +63,7 @@ def center_crop_resize(img, new_center, size): """ CAMERA """ -def intrinsic(): - RES_X = 640 - RES_Y = 480 +def intrinsic(RES_X=640, RES_Y=480): f_mm = 50 # blender default sensor_w_mm = 36 # blender default sensor_h_mm = sensor_w_mm * RES_Y / RES_X @@ -108,3 +106,16 @@ def get_camera(cam, img_size, device): K = torch.tensor(intrinsic()[None, :, :], device=device) cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device)) return cameras, K + + +def filter_info(info, frame_nr): + frame_specific = set(['poses', 'trans']) + general = set(['shape', 'camLoc', 'zrot', 'gender']) + out = dict() + for key in info.keys(): + if key in frame_specific: + out[key] = info[key][:, frame_nr] + elif key in general: + out[key] = info[key] + + return out \ No newline at end of file