fixed normalization, can overfit on single image

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 import math
 import time
@@ -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, filter_info
+from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
 from cloth3d.read import DataReader
 from cloth3d.util import quads2tris, zRotMatrix
 
@@ -65,10 +65,11 @@ class Dataset(Dataset):
         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'))
+        # raw_mask = sio.loadmat(os.path.join(path, frame.garment + '_mask.mat'))['mask'].reshape((-1))
+        # mesh_mask = raw_mask > 0 #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))
+        # 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, zrot, use_trans=True)
         body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
@@ -76,17 +77,22 @@ class Dataset(Dataset):
         body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=self.device)
         body_faces = quads2tris(body_faces)
 
+        joints = torch.tensor(joints[0], device=self.device).to(torch.float32)
+
         if self.random_view:
             x_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
             y_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
             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)
+            # 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)
+            mesh_posed = mesh_posed - joints[0,:][None,:]
+            body = body - joints[0,:][None,:]
+
 
             transform = pytorch3d.transforms.Rotate(view)
             mesh_posed = transform.transform_points(mesh_posed)
@@ -95,11 +101,13 @@ class Dataset(Dataset):
             # 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
+            # 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
+            mesh_posed = mesh_posed + joints[0,:][None,:]
+            body = body + joints[0,:][None,:]
 
         else:
             view = torch.eye(3, device=self.device)
@@ -141,27 +149,32 @@ 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()
+        root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=self.device))[0,:2].round().to(torch.int64)
         
         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)
+        resizer_mask = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
+        resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped, mask_resized, (256, 256))
 
         # 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_garment = resizer_mask(depth_mask_garment[None,None,:,:])[0,0,:,:]
 
+        original = depth_image
+        original_mask = depth_mask
 
-        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,:,:]
+        depth_mask_original = depth_mask
+        depth_mask_croped = cropper(depth_mask)
+        depth_mask = resizer_mask(depth_mask_croped[None,None,:,:])[0,0,:,:]
+        scaled = resizer(cropper(depth_image)[None,None,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,0,:,:]
+        normal_image = resizer(cropper(normal_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
 
         # 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)
+        root_dist = torch.tensor(np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2)), device=self.device).to(torch.float32)
+        # root_dist = torch.tensor(np.abs(root_joint[0][0] - info['camLoc'][0]), device=self.device).to(torch.float32)
         # scaled = root_dist - scaled
         scaled = scaled - root_dist
 
@@ -268,7 +281,7 @@ class Dataset(Dataset):
         # 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).to(self.device)
-        rgb_image = resizer(cropper(rgb_image)[None,:,:,:])[0,:,:,:] * depth_mask[None,:,:]
+        rgb_image = resizer(cropper(rgb_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
 
         # 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)
 
@@ -278,8 +291,7 @@ class Dataset(Dataset):
         data = { 'mesh_posed': mesh_posed,
                  'depth_image': scaled, 
                 #  'depth_image': upscaled, 
-                 'mesh_mask': mesh_mask, 
-                 'depth_mask': depth_mask, #depth_mask,  depth_mask_garment
+                 'depth_mask': depth_mask_garment, #depth_mask,  depth_mask_garment
                  'depth_mask_full': depth_mask, #depth_mask, 
                  'normal_image': normal_image,
                 #  'depth_mask_garment': depth_mask_garment,
@@ -290,10 +302,14 @@ class Dataset(Dataset):
                  'root_dist': root_dist,
                  'scale': scale,
                  'root_joint_pixel': root_joint_pixel,
+                 'joints': joints,
+                 'original': original,
+                 'original_mask': original_mask,
                 }
 
-        self.cache[frame] = data
-        return self.cache[frame]
+        if self.cache is not None:
+            self.cache[frame] = data
+        return data
 
 
 if __name__ == '__main__':
@@ -302,8 +318,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_single.txt', random_view=False)
-    train_dataloader = DataLoader(dataset, batch_size=1, drop_last=True, shuffle=False)
+    dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True)
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader:
         batch

data/preprocess.py

@@ -9,6 +9,7 @@ from tqdm import tqdm
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from utils.data_utils import mask_faces, Frame
 from superSMPL.IO import readPC2, writePC2
+from cloth3d.util import loadInfo
 
 
 def generate_faces(folder, frames):
@@ -71,7 +72,7 @@ def genereated_registered_and_mask(folder, is_skirt = False):
 if __name__ == '__main__':
     # folder = '/home/workspace/local-data/train'
     # folder = '/home/workspace/local-data/valid'
-    folder = 'dataset/CLOTH3D/train'
+    folder = '/home/workspace/data/dataset/cloth3d/train'
 
     # generate_faces(folder)
     # genereated_registered_and_mask(folder)

data/render.py

@@ -146,13 +146,13 @@ class DepthRenderer(nn.Module):
                     raster_settings=raster_settings
                 )
         output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:]))
-        v2i = cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
+        v2i = None #cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
         depth = output.zbuf[0,:,:,0]
         pix_to_face = output.pix_to_face[0,:,:,0]
 
         mask = depth != -1
         # dmin = depth.min()
-        dmean = depth[mask].mean()
+        # dmean = depth[mask].mean()
         # dmax = depth[mask].max()
         # depth[mask] = depth[mask] - dmean
         # depth[~mask] = 0.0

models/depth_estim.py

@@ -108,6 +108,7 @@ class DepthEstimation(nn.Module):
         view_features = self.view_encoder(view)
         # x = self.flatten(x)
         x = self.decoder([torch.concat((x, view_features), axis=1), features])
+        x = torch.tanh(x[0]), x[1]
         # x = self.decoder(x)
         
         return x

test.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 import sys
 
@@ -21,21 +21,63 @@ from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
 ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
 
 
-def plot3d(depth, mask):
+def plot3d(depth, mesh, info, root_dist, scale, mask, root_joint_pixel):
     import kornia
     import plotly.graph_objects as go
-    from data.render import intrinsic
-    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
+    from utils.data_utils import  get_camera
+    import torchvision as tv
+    cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device())
+    mask_orig = mask[:,None,:,:]
+    mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel)
+    depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask
+
+    pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K)
+    pc = pc.detach().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'),
                         ]).show()
 
 
-def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel, device='cuda'):
+def restore(img, scale, center, masks=None, root_dist=0.0):
+    import torchvision as tv
+    import numpy as np
+    from utils.data_utils import resize_image
+    if masks is None:
+        resizer = tv.transforms.Resize((torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
+    else:
+        resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped[:,0,:,:], mask_resized[:,0,:,:], (torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)))
+    
+    cropper = lambda img, off: tv.transforms.functional.crop(img, off[1], off[0], 480, 640)
+
+    if masks is None:
+        out = resizer(img + root_dist)
+    else:
+        out = resizer(img + root_dist, masks[0], masks[1])
+
+    out_croped = out
+    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, out_croped
+    
+
+def depth_to_3d(depth, depth_mask, root_joint_pixel, root_dist, K):
     import kornia
+    # pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    root_joint_image = torch.zeros_like(depth)
+    root_joint_image[:,:,root_joint_pixel[:,1], root_joint_pixel[:,0]] = root_dist
+    root_reprojected = kornia.geometry.depth.depth_to_3d(root_joint_image, K, normalize_points=True)[0,:,root_joint_pixel[:,1], root_joint_pixel[:,0]][:,0]
+    pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    pc = pc[0,:,depth_mask[0,0,:,:]]
+    # pc[2,:] -= root_dist
+    # root_reprojected[2] = root_dist
+    pc = pc - root_reprojected[:,None]
+    return pc
+
+
+def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel, joints, original, original_mask, device='cuda'):
     import plotly.graph_objects as go
     import plotly.express as ex
     from cloth3d.util import zRotMatrix
@@ -44,30 +86,21 @@ def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel
     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())
+    # cameras, K = get_camera(info['camLoc'][0], [640, 480], 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
+
+    # depth = original[None, :, :, :]
+    # mask = original_mask[None, :, :, :]
+
+    mask_orig = mask[:,None,:,:]
+    mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel)
+    depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask
+
+    pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K)
     pc = pc.detach().cpu().numpy()
 
-    mesh = mesh[0]
+    mesh = mesh[0] - joints[0,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)
@@ -103,7 +136,7 @@ def plot_gt(depth, gt, mask):
     pass
 
 
-def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
+def chamfer_loss(depth_image, vertices, depth_mask):
     import kornia
     import pytorch3d
     from data.render import intrinsic
@@ -116,7 +149,7 @@ def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
     pc = [torch.permute(pc[i,:,depth_mask[i]], [1,0]) for i in range(depth_image.shape[0])]
     pc_norm = get_normals_from_depth(depth_image, depth_mask, K)
     pc_norm = [pc_norm[i,depth_mask[i],:] for i in range(pc_norm.shape[0])]
-    mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
+    mesh_gt = [vertices[i, :, :] for i in range(vertices.shape[0])]
     pc_pc = Pointclouds(pc, normals=pc_norm)
     mesh_gt_pc = Pointclouds(mesh_gt)
     mesh_norm = list(torch.unbind(estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)))
@@ -153,6 +186,7 @@ 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
@@ -165,11 +199,18 @@ def normals_reconst_loss(y_pred, normals_gt, mask):
 
 
 def train_step(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'], batch['joints'], batch['original'], batch['original_mask'])
+    # plot3d(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'])
     ### Training
     model.train() # train mode is on by default after construction
 
+    # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1)
+    input = batch['rgb_image']
+    # import plotly.express as ex
+    # ex.imshow(batch['depth_image'][0,:,:].cpu().numpy(), zmin=-1, zmax=5).show()
+
     # 1. Forward pass
-    y_pred, y_pred_normals = model([batch['rgb_image'], batch['view']])
+    y_pred, y_pred_normals = model([input, batch['view']])
 
     # 2. Calculate loss
     mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
@@ -179,10 +220,9 @@ 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 + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst
+    loss = 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()
@@ -201,7 +241,10 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference) 
     # 1. Forward pass
     with torch.inference_mode():
-        pred, pred_normals = model([batch['rgb_image'], batch['view']])
+        # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1)
+        input = batch['rgb_image']
+
+        pred, pred_normals = model([input, batch['view']])
 
         # 2. Calculate the loss
         mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
@@ -211,7 +254,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 + mse_normals #+ chamfer_reconst + mse + normals_reconst
+        loss = mse #+ mse_normals #+ chamfer_reconst + mse + normals_reconst
 
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
@@ -230,9 +273,10 @@ if __name__ == '__main__':
     random_view = False
     use_skip = False
     batch_size = 1
+    use_cache = True
 
-    train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=random_view, device=device)
-    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=random_view, device=device)
+    train_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
+    valid_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
     train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
     valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()
@@ -243,7 +287,7 @@ if __name__ == '__main__':
     loss_fn = nn.MSELoss()
     loss_fn_3d = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
-    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2000, gamma=0.1)
+    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9999, verbose=False)
     # optimizer = optim.SGD(model.parameters(), lr=0.01)
 
     iteration = 0

utils/data_utils.py

@@ -6,7 +6,7 @@ import numpy as np
 import scipy.io as sio
 from math import cos, sin
 from pytorch3d.utils import cameras_from_opencv_projection
-
+import torchvision as tv
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 # from superSMPL.IO import readPC2Frame, readOBJ
@@ -58,6 +58,19 @@ def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
     return normals
 
 
+def resize_image(img, mask_croped, mask_resized, size):
+    resizer = tv.transforms.Resize(size, interpolation=tv.transforms.functional.InterpolationMode.BILINEAR)
+    dilated = kornia.morphology.dilation(img, torch.ones((3,3), device=img.get_device()), engine='convolution')
+    dilated = kornia.morphology.dilation(dilated, torch.ones((3,3), device=img.get_device()), engine='convolution')
+    dilated = torch.permute(dilated, [0,2,3,1])
+    img = torch.permute(img, [0,2,3,1])
+    dilated[mask_croped, :] = img[mask_croped, :]
+    dilated = torch.permute(dilated, [0,3,1,2])
+    out = resizer(dilated)
+    out = out * mask_resized[:,None,:,:]
+    return out
+
+
 def center_crop_resize(img, new_center, size):
     pass