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
+28 -10
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)
depth_mask_garment = resizer(cropper(depth_mask_garment)[None,None,:,:])[0,0,:,:]
# 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 = 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