import os import sys os.environ["CUDA_VISIBLE_DEVICES"] = '6' import math import time import torch import kornia import pytorch3d import torchvision as tv import numpy as np import scipy.io as sio import plotly.express as ex 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, resize_image from cloth3d.read import DataReader from cloth3d.util import quads2tris, zRotMatrix class Dataset(Dataset): def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True): self.folder = folder self.folder_rgb = folder_rgb self.device = device self.random_view = random_view self.transforms = tv.transforms.ToTensor() self.frames = read_frames_file(os.path.join(folder, frames_file)) self.render = DepthRenderer() self.deodr_render : Render = Render(max_depth=10) self.body_faces = np.load('custom_data/smpl_faces.npy') self.reader = DataReader(folder, 'superSMPL') self.cache = dict() if use_cache else None def __len__(self): return len(self.frames) # TODO: make it work with two part outfits def __getitem__(self, idx): frame : Frame = self.frames[idx] if self.cache is not None: if frame in self.cache.keys(): return self.cache[frame] 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) 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, None, zrot) mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=self.device) # 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)) # 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) 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) 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) mesh_posed = mesh_posed - joints[0,:][None,:] body = body - joints[0,:][None,:] 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 mesh_posed = mesh_posed + joints[0,:][None,:] body = body + joints[0,:][None,:] else: view = torch.eye(3, device=self.device) combined = torch.concat((body, mesh_posed)) # 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) 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, faces=combined_faces.cpu().numpy().astype(np.int32)) print(f'deodr time: {time.perf_counter() - start}') 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) # 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) normal_image = torch.permute(normal_image[0], [2,0,1]) print(f'normals time: {time.perf_counter() - start}') depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask top_left = root_joint.astype(np.float32) top_left[:,1] -= 1.25; top_left[:,2] += 1.25 bottom_right = root_joint.astype(np.float32) 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().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_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_mask(depth_mask_garment[None,None,:,:])[0,0,:,:] original = depth_image original_mask = depth_mask 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).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 # 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) # 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' # 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) # 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] # 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() # dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]] # depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] = 10 # diff = np.abs(dist - dist_img) # import plotly.express as ex # import plotly.graph_objects as go # ex.imshow(deodr_depth[:,:,0]).show() # ex.imshow(depth_image.cpu().numpy()).show() # ex.imshow(normal_image[0].cpu().numpy()).show() # trans = np.array(info['trans'], dtype=np.float32)[:, self.frames[idx].frame_nr] # root = info['camLoc'].copy() # root[1:] = 0 # go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'), # go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), # go.Scatter3d(x=[root[0]], y=[root[1]], z=[root[2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), # ]).show() # 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() # v2i_coverage = np.zeros((480, 640), dtype=np.float32) # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1 # 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,:,:,:], 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) info = filter_info(info, frame.frame_nr) # image = self.transforms(image) data = { 'mesh_posed': mesh_posed, 'depth_image': scaled, # 'depth_image': upscaled, '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, # '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, 'joints': joints, 'original': original, 'original_mask': original_mask, } if self.cache is not None: self.cache[frame] = data return data if __name__ == '__main__': from torch.utils.data import DataLoader device = "cuda" if torch.cuda.is_available() else "cpu" # dataset = Dataset('/home/workspace/local-data/train') 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 pass pass