general save
This commit is contained in:
+4
-4
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = '5'
|
||||
|
||||
import math
|
||||
import time
|
||||
@@ -34,7 +34,7 @@ class Dataset(Dataset):
|
||||
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.deodr_render : Render = Render(max_depth=10)
|
||||
self.body_faces = np.load('custom_data/smpl_faces.npy')
|
||||
self.reader = DataReader(folder, 'superSMPL')
|
||||
|
||||
@@ -310,7 +310,7 @@ class Dataset(Dataset):
|
||||
if self.cache is not None:
|
||||
self.cache[frame] = data
|
||||
return data
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
from torch.utils.data import DataLoader
|
||||
@@ -318,7 +318,7 @@ if __name__ == '__main__':
|
||||
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)
|
||||
dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True, device=device)
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
|
||||
for batch in train_dataloader:
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
|
||||
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from time import time
|
||||
from cloth3d.read import DataReader
|
||||
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
|
||||
from superSMPL.smpl_torch import SMPL
|
||||
from data.pose_renderer import PoseRenderer
|
||||
|
||||
|
||||
class Dataset(Dataset):
|
||||
def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
|
||||
self.folder = folder
|
||||
self.device = device
|
||||
self.frames = read_frames_file(os.path.join(folder, frames_file))
|
||||
self.reader = DataReader(folder, 'superSMPL')
|
||||
|
||||
self.cache = dict() if use_cache else None
|
||||
|
||||
self.poses = dict()
|
||||
self.shapes = dict()
|
||||
|
||||
ids = set()
|
||||
for frame in self.frames:
|
||||
ids.add(frame.id)
|
||||
|
||||
for id in tqdm(ids):
|
||||
info = self.reader.read_info(id)
|
||||
self.poses[id] = info['poses'].astype(np.float32)
|
||||
self.shapes[id] = np.concatenate((info['shape'], [info['gender']])).astype(np.float32)
|
||||
|
||||
self.smpl = SMPL(path='superSMPL', device=device)
|
||||
self.renderer = PoseRenderer()
|
||||
# self.colors = np.load('custom_data/pose_colors.npy')
|
||||
|
||||
|
||||
def _render_body(self, body_to_render: torch.Tensor):
|
||||
# swap the y,z axes of the body_to_render
|
||||
body_to_render = body_to_render[0,:,[0,2,1]]
|
||||
# and rotate it 90 degress along the z axis
|
||||
body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
|
||||
|
||||
out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.get_device()))
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(self.frames)
|
||||
|
||||
|
||||
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]
|
||||
|
||||
pose = torch.tensor(self.poses[frame.id][:,frame.frame_nr], device=self.device)
|
||||
shape = torch.tensor(self.shapes[frame.id], device=self.device)
|
||||
pose[:3] = 0
|
||||
pose_out = pose.clone()
|
||||
pose_out[1] = np.random.uniform(-np.pi, np.pi)
|
||||
# pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
|
||||
# generate random number between -np.pi, np.pi
|
||||
|
||||
# start = time()
|
||||
|
||||
body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
|
||||
body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
|
||||
|
||||
render = self._render_body(body[0])[0,:,:,:3]
|
||||
render_out = self._render_body(body_out[0])[0,:,:,:3]
|
||||
|
||||
render = render.permute(2,0,1)
|
||||
render_out = render_out.permute(2,0,1)
|
||||
|
||||
# import plotly.express as ex
|
||||
# ex.imshow(render[0,:,:,:3].cpu().numpy()).show()
|
||||
# ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show()
|
||||
|
||||
mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32)
|
||||
mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)
|
||||
|
||||
render = render * mask
|
||||
render_out = render_out * mask_out
|
||||
|
||||
# print(f'one frame took {time() - start} seconds')
|
||||
|
||||
data = {
|
||||
'pose': render,
|
||||
'shape': shape,
|
||||
'pose_out': render_out,
|
||||
# 'body': body,
|
||||
# 'body_out': body_out,
|
||||
# 'render': render,
|
||||
# 'render_out': render_out,
|
||||
'mask': mask,
|
||||
'mask_out': mask_out,
|
||||
}
|
||||
|
||||
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('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
|
||||
for batch in train_dataloader:
|
||||
batch
|
||||
pass
|
||||
|
||||
pass
|
||||
+36
-5
@@ -3,10 +3,10 @@ import sys
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from deodr import Camera, Scene3D
|
||||
from deodr.triangulated_mesh import TriMesh
|
||||
from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
|
||||
from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
|
||||
# from deodr import Camera, Scene3D
|
||||
# from deodr.triangulated_mesh import TriMesh
|
||||
# from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
|
||||
# from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from cloth3d.util import proj
|
||||
@@ -106,7 +106,8 @@ from pytorch3d.renderer import (
|
||||
MeshRasterizer,
|
||||
SoftPhongShader,
|
||||
TexturesUV,
|
||||
TexturesVertex
|
||||
TexturesVertex,
|
||||
HardPhongShader,
|
||||
)
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.utils import cameras_from_opencv_projection
|
||||
@@ -162,3 +163,33 @@ class DepthRenderer(nn.Module):
|
||||
# ex.imshow(depth.cpu().numpy()).show()
|
||||
|
||||
return depth, mask, v2i, K, cameras, pix_to_face
|
||||
|
||||
|
||||
class PoseRenderer(nn.Module):
|
||||
def __init__(self, device='cuda') -> None:
|
||||
super().__init__()
|
||||
colors = np.load('custom_data/pose_colors.npy')
|
||||
self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
|
||||
self.cam = torch.tensor([0.0, 0.0, 1.0], device=device)
|
||||
|
||||
|
||||
def forward(self, verts: torch.Tensor, faces: torch.Tensor):
|
||||
cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
|
||||
|
||||
raster_settings = RasterizationSettings(
|
||||
image_size=(480, 640),
|
||||
blur_radius=0.0,
|
||||
faces_per_pixel=1,
|
||||
bin_size=0
|
||||
)
|
||||
|
||||
renderer = MeshRenderer(
|
||||
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
|
||||
shader=HardPhongShader(device=verts.get_device(), cameras=cameras)
|
||||
)
|
||||
# create a pytorch3d vertex texture from the colors
|
||||
|
||||
output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
|
||||
|
||||
|
||||
return output
|
||||
|
||||
Reference in New Issue
Block a user