fixed normalization, can overfit on single image
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user