125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
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 |