138 lines
5.8 KiB
Python
138 lines
5.8 KiB
Python
import os
|
|
import sys
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
|
|
|
import math
|
|
import time
|
|
import torch
|
|
import kornia
|
|
import torchvision as tv
|
|
import numpy as np
|
|
import scipy.io as sio
|
|
import plotly.express as ex
|
|
|
|
from torch.utils.data import Dataset
|
|
|
|
from render import Render, DepthRenderer
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
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, frames_file='frames_all.txt'):
|
|
self.folder = folder
|
|
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')
|
|
|
|
|
|
def __len__(self):
|
|
return len(self.frames)
|
|
|
|
|
|
# TODO: make it work with two part outfits
|
|
def __getitem__(self, idx):
|
|
frame : Frame = self.frames[idx]
|
|
path = os.path.join(self.folder, frame.id)
|
|
# 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)
|
|
|
|
mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
|
|
# TODO: what's the difference between masking and reading .obj file
|
|
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 = self.reader.read_human(frame.id, frame.frame_nr)
|
|
body_faces = quads2tris(body_faces)
|
|
|
|
info = self.reader.read_info(frame.id)
|
|
combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to('cuda')
|
|
combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda')
|
|
|
|
idx_frame = 36000
|
|
x = combined[idx_frame, :].cpu().detach().numpy()
|
|
y = np.array(info['camLoc'], dtype=np.float32)
|
|
dist = np.sqrt(np.sum((x - y) ** 2))
|
|
|
|
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 = self.render(combined, combined_faces, info['camLoc'])
|
|
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,:,:], depth_mask[None,:,:], K)
|
|
print(f'normals time: {time.perf_counter() - start}')
|
|
|
|
root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2]
|
|
|
|
# Testing
|
|
# pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True)
|
|
# pc = pc[0,:,depth_mask].cpu().numpy()
|
|
|
|
# depth_image = depth_image.cpu().numpy()
|
|
# dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]]
|
|
# depth_image[v2i[idx, 1], v2i[idx, 0]] = 100
|
|
combined = combined.cpu().numpy()
|
|
|
|
# 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).show()
|
|
# ex.imshow(normal_image[0].cpu().numpy()).show()
|
|
|
|
# trans = np.array(info['trans'], dtype=np.float32)[:, self.frames[idx].frame_nr]
|
|
trans = root_joint[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=[trans[0]], y=[trans[1]], z=[trans[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()
|
|
|
|
# image = self.transforms(image)
|
|
return { 'mesh_posed': mesh_posed,
|
|
'depth_image': depth_image,
|
|
'mesh_mask': mesh_mask,
|
|
'depth_mask': depth_mask,
|
|
'normal_image': normal_image
|
|
}
|
|
|
|
|
|
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')
|
|
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
|
|
|
for batch in train_dataloader:
|
|
batch
|
|
pass
|
|
|
|
pass
|