change feature size based on 256x256 input

data/dataset.py

@@ -88,9 +88,9 @@ class Dataset(Dataset):
 
             transform = pytorch3d.transforms.Rotate(view)
             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_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)
@@ -127,60 +127,79 @@ class Dataset(Dataset):
         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)
+        normal_image = torch.permute(normal_image[0], [2,0,1])
         print(f'normals time: {time.perf_counter() - start}')
 
         depth_mask_body = (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
@@ -235,21 +254,22 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # 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)
 
 
         # image = self.transforms(image)
         data = { 'mesh_posed': mesh_posed,
-                 'depth_image': depth_image, 
+                 'depth_image': scaled, 
                 #  'depth_image': upscaled, 
                  'mesh_mask': mesh_mask, 
                  'depth_mask': depth_mask, 
                  'normal_image': normal_image,
                  '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,
                  'info': info,
                 }
@@ -264,8 +284,8 @@ 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', 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:
         batch

models/convnext.py

@@ -241,12 +241,25 @@ class ConvNeXtReverseBase2(nn.Module):
                     ).to(device)
                 )
 
+        for cnf in block_setting[-1:]:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                stage.append(block(cnf.input_channels, layer_scale, 0))
+            self.stage_normals = nn.Sequential(*stage).to(device)
+
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
                     1,
                     kernel_size=3,
                     padding='same',
                     ).to(device)
+        self.last_normals = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
         # self.last = nn.Sequential(
         #                 Permute([0, 2, 3, 1]),
         #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
@@ -272,10 +285,15 @@ class ConvNeXtReverseBase2(nn.Module):
             x = self.upsamples[i](x)
         
         x = self.stages[-2](x)
-        x = self.upsamples[-1](x)
+        f = self.upsamples[-1](x)
-        x = self.stages[-1](x)
+
-        x = self.last(x)
+        x = self.stages[-1](f)
-        return x
+        depth = self.last(x)
+
+        x = self.stage_normals(f)
+        normals = self.last_normals(x)
+        
+        return depth, normals
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)

models/depth_estim.py

@@ -26,10 +26,11 @@ class ViewEncoder(nn.Module):
         x = self.act1(x)
         x = self.lin2(x)
         x = self.act2(x)
-        x = torch.ones((x.shape[0], self.features, 15, 20), device=x.get_device()) * x[:,:,None,None]
+        x = torch.ones((x.shape[0], self.features, 8, 8), device=x.get_device()) * x[:,:,None,None]
 
         return x
 
+
 from torchvision.ops.misc import Conv2dNormActivation, Permute
 
 class DumbReconstructor(nn.Module):

test.py

@@ -34,30 +34,17 @@ def plot3d(depth, mask):
 
 
 # temp
-def plot_gt(y_pred, batch):
+def plot_gt(depth, gt, mask):
     import kornia
     import plotly.graph_objects as go
-    import pytorch3d
     from data.render import intrinsic
-
     K = torch.tensor(intrinsic()[None, :, :], device='cuda')
-    pc = kornia.geometry.depth.depth_to_3d(y_pred - batch['info']['camLoc'][:,0][:,None,None,None].cuda(), K, normalize_points=True)
+    pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
-    pc = pc[0,:,batch['depth_mask'][0]].detach().cpu().numpy()
+    pc = pc[0,:,mask[0]].detach().cpu().numpy()
-    m = batch['mesh_posed'][0, batch['mesh_mask'][0], :]
+    pc_gt = kornia.geometry.depth.depth_to_3d(gt[:,None,:,:], K, normalize_points=True)
-    view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((torch.tensor([torch.pi/2], device='cuda'),
+    pc_gt = pc_gt[0,:,mask[0]].detach().cpu().numpy()
-                                                                 torch.tensor([0.0], device='cuda'),
-                                                                 torch.tensor([0.0], device='cuda'))))
-    transform = pytorch3d.transforms.Rotate(view)
-    m = transform.transform_points(m)
-    view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((torch.tensor([0.0], device='cuda'),
-                                                                 torch.tensor([torch.pi/2], device='cuda'),
-                                                                 torch.tensor([0.0], device='cuda'))))
-    transform = pytorch3d.transforms.Rotate(view)
-    m = transform.transform_points(m)
-    m = m.cpu().numpy()
-
     go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
-    go.Scatter3d(x=m[:,0], y=m[:,1], z=m[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                        go.Scatter3d(x=pc_gt[0,:], y=pc_gt[1,:], z=pc_gt[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
                         ]).show()
 
     pass
@@ -113,21 +100,33 @@ def depth_reconst_loss(y_pred, y_true, mask):
     chamfer, _ = chamfer_pytorch(pred_pc_pc, gt_pc_pc)
     return chamfer
 
+def normals_reconst_loss(y_pred, normals_gt, mask):
+    import kornia
+    from data.render import intrinsic
+    from utils.data_utils import  get_normals_from_depth
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    normal_image = torch.permute(get_normals_from_depth(y_pred, mask, K), [0,3,1,2])
+
+    mse = loss_fn(normal_image * mask[:,None,:,:], normals_gt * mask[:,None,:,:])
+    return mse
+
 
 def train_step(batch):
     ### Training
     model.train() # train mode is on by default after construction
 
     # 1. Forward pass
-    y_pred = model([batch['rgb_image'], batch['view']])
+    y_pred, y_pred_normals = model([batch['rgb_image'], batch['view']])
 
     # 2. Calculate loss
     mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+    mse_normals = loss_fn(y_pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:])
     # chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
     msssim = ms_ssim_module(y_pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5)
-    chamfer_reconst = depth_reconst_loss(y_pred, batch['depth_image'][:,None,:,:],  batch['depth_mask'])
+    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 + chamfer_reconst + mse
+    loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst
 
     # plot_gt(y_pred, batch)
 
@@ -148,15 +147,17 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference) 
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model([batch['rgb_image'], batch['view']])
+        pred, pred_normals = model([batch['rgb_image'], batch['view']])
 
         # 2. Calculate the loss
         mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+        mse_normals = loss_fn(pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:])
         # chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
         msssim = ms_ssim_module(pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5)
         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 + chamfer_reconst + mse
+        loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + normals_reconst
 
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
@@ -187,7 +188,8 @@ if __name__ == '__main__':
 
     loss_fn = nn.MSELoss()
     loss_fn_3d = nn.MSELoss()
-    optimizer = optim.Adam(model.parameters(), lr=0.01)
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2000, gamma=0.1)
     # optimizer = optim.SGD(model.parameters(), lr=0.01)
 
     iteration = 0
@@ -197,19 +199,22 @@ if __name__ == '__main__':
             train_loss = train_step(batch)
             train_accumulator.update(train_loss)
 
-            if iteration % 1000 == 0:
+            lr_scheduler.step()
-                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
-                train_accumulator.reset()
-
             iteration += 1
 
+        if iteration % 1000 == 0:
+            print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
+            train_accumulator.reset()
+
+        # TODO: change when batchsize != 1
         if iteration % 100 == 0:
             for batch in valid_dataloader:
                 valid_loss = valid_step(batch)
                 valid_accumulator.update(valid_loss)
 
             print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
-        valid_accumulator.reset()
+            valid_accumulator.reset()
+