luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

save

Browse Source
This commit is contained in:
luni
2023-01-23 10:17:53 +00:00
parent f4341c5e6c
commit ab78dab165
4 changed files with 112 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+27 -9
View File
@@ -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
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info
from cloth3d.read import DataReader
from cloth3d.util import quads2tris, zRotMatrix
@@ -54,6 +54,8 @@ class Dataset(Dataset):
path = os.path.join(self.folder, frame.id)
info = self.reader.read_info(frame.id)
info['zrot'] = torch.pi / 2 + torch.pi/4
zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
trans = info['trans'][:,frame.frame_nr].astype(np.float32)
# mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
@@ -103,7 +105,7 @@ class Dataset(Dataset):
view = torch.eye(3, device=self.device)
combined = torch.concat((body, mesh_posed[mesh_mask, :]))
combined_no_trans = torch.concat((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))
# 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)
@@ -111,7 +113,7 @@ class Dataset(Dataset):
start = time.perf_counter()
# depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
# depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32))
# deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32))
# deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined, faces=combined_faces.cpu().numpy().astype(np.int32))
print(f'deodr time: {time.perf_counter() - start}')
@@ -119,7 +121,7 @@ class Dataset(Dataset):
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)
# 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}')
@@ -139,18 +141,29 @@ 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()
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)
# 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(cropper(depth_mask_garment)[None,None,:,:])[0,0,:,:]
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,:,:]
root_dist = np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2))
scaled = root_dist - scaled
# 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)
# scaled = root_dist - scaled
scaled = scaled - root_dist
# 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]
@@ -259,19 +272,24 @@ class Dataset(Dataset):
# 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)
info = filter_info(info, frame.frame_nr)
# image = self.transforms(image)
data = { 'mesh_posed': mesh_posed,
'depth_image': scaled,
# 'depth_image': upscaled,
'mesh_mask': mesh_mask,
'depth_mask': depth_mask,
'depth_mask': depth_mask, #depth_mask, depth_mask_garment
'depth_mask_full': depth_mask, #depth_mask,
'normal_image': normal_image,
'depth_mask_garment': depth_mask_garment,
# 'depth_mask_garment': depth_mask_garment,
# 'depth_mask_body': depth_mask_body,
'rgb_image': rgb_image,
'view': view,
'info': info,
'root_dist': root_dist,
'scale': scale,
'root_joint_pixel': root_joint_pixel,
}
self.cache[frame] = data
speed.py
+13
View File
@@ -0,0 +1,13 @@
import numpy as np
import time
if __name__ == '__main__':
start = time.perf_counter()
a, b = np.random.rand(1024, 1024), np.random.rand(1024, 1024)
print(f'init time: {time.perf_counter() - start}')
start = time.perf_counter()
c = np.dot(a, b)
print(f'dot time: {time.perf_counter() - start}')
test.py
+57 -3
View File
@@ -25,14 +25,67 @@ def plot3d(depth, mask):
import kornia
import plotly.graph_objects as go
from data.render import intrinsic
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
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
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'):
import kornia
import plotly.graph_objects as go
import plotly.express as ex
from cloth3d.util import zRotMatrix
import pytorch3d
import numpy as np
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())
# 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
pc = pc.detach().cpu().numpy()
mesh = mesh[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)
# zRot = torch.tensor(zRotMatrix(info['zrot']), device=device)
# mesh = torch.matmul(zRot, mesh.T).T
# view = pytorch3d.transforms.axis_angle_to_matrix(torch.tensor([0.0, 0.0, torch.pi/2], device=device))
# transform = pytorch3d.transforms.Rotate(view)
# mesh = transform.transform_points(mesh)
mesh = mesh.detach().cpu().numpy()
go.Figure(data = [go.Scatter3d(x=pc[2,:], y=pc[0,:], z=-pc[1,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
go.Scatter3d(x=-mesh[:,0], y=mesh[:,1], z=mesh[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
]).show()
pass
# temp
def plot_gt(depth, gt, mask):
import kornia
@@ -126,9 +179,10 @@ 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_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst
loss = 1 - msssim + 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()
@@ -157,7 +211,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_normals #+ chamfer_reconst + mse + normals_reconst
loss = 1 - msssim + mse + mse_normals #+ chamfer_reconst + mse + normals_reconst
return { 'valid/loss': loss.detach().cpu().numpy() }
utils/data_utils.py
+14 -3
View File
@@ -63,9 +63,7 @@ def center_crop_resize(img, new_center, size):
""" CAMERA """
def intrinsic():
RES_X = 640
RES_Y = 480
def intrinsic(RES_X=640, RES_Y=480):
f_mm = 50 # blender default
sensor_w_mm = 36 # blender default
sensor_h_mm = sensor_w_mm * RES_Y / RES_X
@@ -108,3 +106,16 @@ def get_camera(cam, img_size, device):
K = torch.tensor(intrinsic()[None, :, :], device=device)
cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device))
return cameras, K
def filter_info(info, frame_nr):
frame_specific = set(['poses', 'trans'])
general = set(['shape', 'camLoc', 'zrot', 'gender'])
out = dict()
for key in info.keys():
if key in frame_specific:
out[key] = info[key][:, frame_nr]
elif key in general:
out[key] = info[key]
return out