diff --git a/data/dataset.py b/data/dataset.py index 222df18..f0398a1 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -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: diff --git a/models/convnext.py b/models/convnext.py index d18bb2c..14798cc 100644 --- a/models/convnext.py +++ b/models/convnext.py @@ -244,14 +244,14 @@ class ConvNeXtReverseBase2(nn.Module): self.last = nn.Conv2d( block_setting[-1].input_channels, 1, - kernel_size=7, - padding=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), # Permute([0, 3, 1, 2]), - # ) + # ).to(device) for m in self.modules(): diff --git a/models/depth_estim.py b/models/depth_estim.py index 9d20610..f64753a 100644 --- a/models/depth_estim.py +++ b/models/depth_estim.py @@ -16,7 +16,7 @@ class ViewEncoder(nn.Module): self.flatten = nn.Flatten() self.lin1 = nn.LazyLinear(features).to(device) self.act1 = nn.GELU().to(device) - self.lin2 = nn.Linear(features, features * 15 * 20).to(device) + self.lin2 = nn.Linear(features, features).to(device) self.act2 = nn.Tanh().to(device) @@ -26,10 +26,51 @@ class ViewEncoder(nn.Module): x = self.act1(x) x = self.lin2(x) x = self.act2(x) - x = x.view(x.shape[0],self.features,15,20) + x = torch.ones((x.shape[0], self.features, 15, 20), device=x.get_device()) * x[:,:,None,None] return x +from torchvision.ops.misc import Conv2dNormActivation, Permute + +class DumbReconstructor(nn.Module): + def __init__(self, view_features=128, device='cuda'): + super().__init__() + + # self.model = torch.nn.Sequential( + # nn.ConvTranspose2d(768 + view_features, 128, 7, 4, 3), + # nn.BatchNorm2d(128), + # nn.ConvTranspose2d(128, 128, 7, 4, 3), + # nn.Conv2d(128, 64, 7, 1, 'same'), + # nn.BatchNorm2d(64), + # nn.ConvTranspose2d(64, 32, 7, 4, 3), + # nn.Conv2d(32, 1, 5, 1, 'same') + # ).to(device) + + self.last = nn.Sequential( + nn.Linear(in_features=1024, out_features=480*640, bias=True), + ).to(device) + + def forward(self, input) -> torch.Tensor: + x = input + x = self.last(x).view((1,1,480,640)) + return x + + +class DumbEncoder(nn.Module): + def __init__(self, view_features=128, device='cuda'): + super().__init__() + + self.last = nn.Sequential( + nn.Flatten(), + nn.Linear(in_features=3*480*640, out_features=1024, bias=True), + ).to(device) + + def forward(self, input) -> torch.Tensor: + x = input + x = self.last(x) + return x, None + + class DepthEstimation(nn.Module): def __init__(self, view_features=128, device='cuda', use_skip=False): super().__init__() @@ -52,8 +93,10 @@ class DepthEstimation(nn.Module): ] self.convnext = ConvNeXt(block_setting, device=device) + # self.convnext = DumbEncoder() self.flatten = nn.Flatten() self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device) + # self.decoder = DumbReconstructor() self.view_encoder = ViewEncoder(view_features) @@ -63,6 +106,7 @@ class DepthEstimation(nn.Module): x, features = self.convnext(image) view_features = self.view_encoder(view) # x = self.flatten(x) - x = self.decoder([torch.concat((x, view_features), axis=1),features]) + x = self.decoder([torch.concat((x, view_features), axis=1), features]) + # x = self.decoder(x) return x diff --git a/requirements.txt b/requirements.txt index 6d28766..fa8a232 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ plotly scipy deodr -pandas \ No newline at end of file +pandas +pytorch_msssim \ No newline at end of file diff --git a/test.py b/test.py index eaf973c..58ce26e 100644 --- a/test.py +++ b/test.py @@ -17,7 +17,8 @@ from utils.utils import Accumulator from models.depth_estim import DepthEstimation from models.convnext import ConvNeXt - +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): @@ -32,6 +33,36 @@ def plot3d(depth, mask): ]).show() +# temp +def plot_gt(y_pred, batch): + 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 = pc[0,:,batch['depth_mask'][0]].detach().cpu().numpy() + m = batch['mesh_posed'][0, batch['mesh_mask'][0], :] + view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((torch.tensor([torch.pi/2], device='cuda'), + 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'), + ]).show() + + pass + + def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask): import kornia import pytorch3d @@ -56,15 +87,32 @@ def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask): return chamfer -def normal_loss(depth_image, depth_mask, normal_image_gt): +# def normal_loss(depth_image, depth_mask, normal_image_gt): +# import kornia +# import pytorch3d +# from data.render import intrinsic +# K = torch.tensor(intrinsic()[None, :, :], device='cuda') +# pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True) +# pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])]) +# normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K) +# pass + +def depth_reconst_loss(y_pred, y_true, mask): import kornia - import pytorch3d from data.render import intrinsic + from pytorch3d.structures import Pointclouds + from pytorch3d.loss.chamfer import chamfer_distance as chamfer_pytorch + # from utils.pytorch3d_chamfer import chamfer_distance K = torch.tensor(intrinsic()[None, :, :], device='cuda') - pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True) - pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])]) - normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K) - pass + pred_pc = kornia.geometry.depth.depth_to_3d(y_pred, K, normalize_points=True) + pred_pc = [torch.permute(pred_pc[i, :, mask[i]], [1,0]) for i in range(y_pred.shape[0])] + pred_pc_pc = Pointclouds(pred_pc) + gt_pc = kornia.geometry.depth.depth_to_3d(y_true, K, normalize_points=True) + gt_pc = [torch.permute(gt_pc[i, :, mask[i]], [1,0]) for i in range(y_true.shape[0])] + gt_pc_pc = Pointclouds(gt_pc) + chamfer, _ = chamfer_pytorch(pred_pc_pc, gt_pc_pc) + return chamfer + def train_step(batch): ### Training @@ -75,8 +123,13 @@ def train_step(batch): # 2. Calculate loss mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) - chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) - loss = mse + 1e-1 * (chamfer + normal) + # 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']) + # loss = mse #+ 1e-1 * (chamfer + normal) + loss = 1 - msssim + chamfer_reconst + mse + + # plot_gt(y_pred, batch) # 3. Zero grad optimizer optimizer.zero_grad() @@ -99,8 +152,12 @@ def valid_step(batch): # 2. Calculate the loss mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) - chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) - loss = mse + 1e-1 * (chamfer + normal) + # 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']) + # loss = mse #+ 1e-1 * (chamfer + normal) + loss = 1 - msssim + chamfer_reconst + mse + return { 'valid/loss': loss.detach().cpu().numpy() } @@ -117,10 +174,10 @@ if __name__ == '__main__': random_view = False use_skip = False - batch_size = 2 + batch_size = 1 - train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device) - valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device) + 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_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() @@ -129,7 +186,9 @@ if __name__ == '__main__': model = DepthEstimation(device=device, use_skip=use_skip) loss_fn = nn.MSELoss() - optimizer = optim.Adam(model.parameters(), lr=0.001) + loss_fn_3d = nn.MSELoss() + optimizer = optim.Adam(model.parameters(), lr=0.01) + # optimizer = optim.SGD(model.parameters(), lr=0.01) iteration = 0 while True: @@ -144,11 +203,12 @@ if __name__ == '__main__': iteration += 1 - for batch in valid_dataloader: - valid_loss = valid_step(batch) - valid_accumulator.update(valid_loss) + 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']}") + print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}") valid_accumulator.reset()