import os import sys os.environ["CUDA_VISIBLE_DEVICES"] = '0' 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 from cloth3d.read import DataReader from cloth3d.util import quads2tris 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) zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0] trans = info['trans'][:,frame.frame_nr] # 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_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 = torch.tensor(body.astype(np.float32), device=self.device) body_faces = quads2tris(body_faces) 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))) transform = pytorch3d.transforms.Rotate(view) mesh_posed = transform.transform_points(mesh_posed) body = transform.transform_points(body) 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_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) 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) 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 # 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) # Testing # 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)) # 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) # 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) # image = self.transforms(image) data = { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mesh_mask': mesh_mask, 'depth_mask': depth_mask, 'normal_image': normal_image, 'depth_mask_garment': depth_mask_garment, 'depth_mask_body': depth_mask_body, 'rgb_image': rgb_image.to(self.device), 'view': view, } self.cache[frame] = data return self.cache[frame] 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('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt') train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: batch pass pass