general save

data/pose_dataset.py

@@ -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