save

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() }