backup
This commit is contained in:
+91
-65
@@ -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:
|
||||
|
||||
+3
-3
@@ -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():
|
||||
|
||||
+47
-3
@@ -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
|
||||
|
||||
@@ -2,3 +2,4 @@ plotly
|
||||
scipy
|
||||
deodr
|
||||
pandas
|
||||
pytorch_msssim
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user