diff --git a/data/dataset.py b/data/dataset.py index b81e006..8821fa3 100644 --- a/data/dataset.py +++ b/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 diff --git a/data/preprocess.py b/data/preprocess.py index ee3bea1..80ddb01 100644 --- a/data/preprocess.py +++ b/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) diff --git a/data/render.py b/data/render.py index 46c09d1..69c73fa 100644 --- a/data/render.py +++ b/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 diff --git a/models/depth_estim.py b/models/depth_estim.py index dc7080a..2b9ddab 100644 --- a/models/depth_estim.py +++ b/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 diff --git a/test.py b/test.py index ad2b19b..d57de5a 100644 --- a/test.py +++ b/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 diff --git a/utils/data_utils.py b/utils/data_utils.py index 8644e35..8ab55fe 100644 --- a/utils/data_utils.py +++ b/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