general save
This commit is contained in:
@@ -0,0 +1,149 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '5'
|
||||
|
||||
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
|
||||
import kornia
|
||||
import plotly.graph_objects as go
|
||||
from data.render import intrinsic
|
||||
from simpleicp import SimpleICP, PointCloud
|
||||
|
||||
|
||||
def depth_to_3d(depth, depth_mask):
|
||||
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
||||
pc_full = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
|
||||
pc = [pc_full[i,:,depth_mask[i]] for i in range(len(depth_mask))]
|
||||
# x = torch.sum(torch.stack(depth_mask), dim=0).cpu().numpy() == len(depth_mask)
|
||||
# choice = [0,0]
|
||||
# while x[choice[0], choice[1]] == False:
|
||||
# choice = (np.random.random(2) * [480, 640]).astype(int)
|
||||
# for i in range(len(pc)):
|
||||
# pc[i] = pc[i] - pc_full[i][:,choice[0], choice[1]][:,None]
|
||||
return pc
|
||||
|
||||
|
||||
def perspective_test(folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, fix_camera=True):
|
||||
transforms = tv.transforms.ToTensor()
|
||||
frames = read_frames_file(os.path.join(folder, frames_file))
|
||||
render = DepthRenderer()
|
||||
body_faces = np.load('custom_data/smpl_faces.npy')
|
||||
reader = DataReader(folder, 'superSMPL')
|
||||
|
||||
frame : Frame = frames[0]
|
||||
|
||||
path = os.path.join(folder, frame.id)
|
||||
info = 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 = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
|
||||
mesh_posed = torch.tensor(mesh_posed, device=device)
|
||||
# mesh_posed_no_trans = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
|
||||
# mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=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(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 = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
|
||||
body_no_trans, _, _, _, _ = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
|
||||
body = torch.tensor(body.astype(np.float32), device=device)
|
||||
body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=device)
|
||||
body_faces = quads2tris(body_faces)
|
||||
|
||||
joints = torch.tensor(joints[0], device=device).to(torch.float32)
|
||||
|
||||
offsets = [[1, -0.9, 0], [1, 0.9, 0], [1, 0, 0.4], [1, 0, -0.4], [1,-0.9,0.4], [1, 0.9, -0.4]]
|
||||
# x: - = closer, + = farther (+-0.5)
|
||||
# y: - = to the right, + = to the left (+-1)
|
||||
# z: - = up, + = down (+-0.5)
|
||||
|
||||
combineds = []
|
||||
for offset in offsets:
|
||||
combineds.append(torch.concat((body, mesh_posed)) - torch.tensor(offset, device=body.get_device())[None,:])
|
||||
# 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(device)
|
||||
|
||||
depth_images = []
|
||||
depth_masks = []
|
||||
for i, combined in enumerate(combineds):
|
||||
start = time.perf_counter()
|
||||
cam = info['camLoc'].copy()
|
||||
if not fix_camera:
|
||||
cam = cam - offsets[i]
|
||||
depth_image, depth_mask, _, _, _, _ = render(combined, combined_faces, cam)
|
||||
print(f'pytorch3d time: {time.perf_counter() - start}')
|
||||
depth_images.append(depth_image)
|
||||
depth_masks.append(depth_mask)
|
||||
|
||||
# depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
|
||||
# depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask
|
||||
|
||||
# rgb_image = tv.io.read_image(os.path.join(folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32).to(device)
|
||||
|
||||
all_depth = torch.stack(depth_images)
|
||||
pointclouds = depth_to_3d(all_depth[:,None,:,:], depth_masks)
|
||||
|
||||
all_depth_merged = torch.sum(all_depth, dim=0)
|
||||
ex.imshow(all_depth_merged.cpu()).show()
|
||||
|
||||
mesh_data = []
|
||||
for pc in pointclouds:
|
||||
pc = pc.cpu().numpy()
|
||||
mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)
|
||||
|
||||
go.Figure(data = mesh_data).show()
|
||||
|
||||
pc_0 = pointclouds[0] - torch.mean(pointclouds[0], dim=1)[:,None]
|
||||
pc_fix = PointCloud(torch.permute(pc_0,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
|
||||
pc_registered = [np.transpose(pc_0.cpu().numpy(), (1,0))]
|
||||
for pc in pointclouds[1:]:
|
||||
pc = pc - torch.mean(pc, dim=1)[:,None]
|
||||
# pc_reg = torch.permute(pc,[1,0]).cpu().numpy()
|
||||
pc_mov = PointCloud(torch.permute(pc,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
|
||||
icp = SimpleICP()
|
||||
icp.add_point_clouds(pc_fix, pc_mov)
|
||||
_, pc_reg, _, _ = icp.run(max_iterations=50, correspondences=10000, min_change=0.5, neighbors=50)
|
||||
pc_registered.append(pc_reg)
|
||||
# pass
|
||||
|
||||
mesh_data = []
|
||||
for pc in pc_registered:
|
||||
pc = np.transpose(pc, (1,0))
|
||||
mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)
|
||||
|
||||
go.Figure(data = mesh_data).show()
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
fix_camera = True
|
||||
perspective_test('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', fix_camera=fix_camera, device=device)
|
||||
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user