downscale cape

data/downscale_cape.py

@@ -0,0 +1,207 @@
+import argparse
+import os
+import tqdm
+import cv2
+import scipy.interpolate as spi
+import numpy as np
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+
+def get_img_idxs():
+    img_idxs = dict()
+    img_coords = dict()
+    uv_masks = dict()
+    uv_mask = np.load('custom_data/garment_uv_mask_full.npy')
+    coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0))
+    img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0))
+    img_idxs[512] = img_idx
+    img_coords[512] = coords
+    uv_masks[512] = uv_mask.astype(np.int32)
+    for res in [32, 64, 128, 256]:
+        down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy()
+        down = np.squeeze(down.astype(bool))
+        coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
+        img_idx = np.transpose(np.stack(np.nonzero((down))), (1,0))
+        img_idxs[res] = img_idx
+        img_coords[res] = coords
+        uv_masks[res] = down.astype(np.int32)
+
+
+    uv_masks_body = dict()
+    uv_mask_body = np.load('custom_utils/mask_256x512.npy')
+    uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
+    uv_masks_body[512] = uv_mask_body
+    for res in [32, 64, 128, 256]:
+        down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy()
+        down = np.squeeze(down.astype(bool))
+        uv_masks_body[res] = down
+
+    return img_idxs, img_coords, uv_masks, uv_masks_body
+
+
+def fill(img, mask):
+    coords = img_coords[img.shape[0]]
+    img_idx = img_idxs[img.shape[0]]
+    # coords = np.transpose(np.stack(np.nonzero((1 - uv_masks[img.shape[0]]))), (1,0))
+    # img_idx = np.transpose(np.stack(np.nonzero((uv_masks[img.shape[0]]))), (1,0))
+    img_data = img[img_idx[:,0], img_idx[:,1], :]
+
+    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+    img_2 = interpolator(coords)
+    res = np.zeros((512,256,3), dtype=np.float32)
+    res[coords[:,0], coords[:,1], :] = img_2
+    res[img_idx[:,0], img_idx[:,1], :] = img_data
+
+    return res
+
+
+def fill_label(img, mask):
+    coords = img_coords[img.shape[0]]
+    img_idx = img_idxs[img.shape[0]]
+    img_data = img[img_idx[:,0], img_idx[:,1]]
+
+    interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
+    img_2 = interpolator(coords)
+    res = np.zeros((512,256), dtype=np.float32)
+    res[coords[:,0], coords[:,1]] = img_2
+    res[img_idx[:,0], img_idx[:,1]] = img_data
+
+    return res
+
+
+def generate_downscaled_uv_map(folder, id, size):
+    img = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); img = img['uv_mask'].astype(np.float32)
+    down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+    down = np.squeeze(down.numpy().astype(bool))
+    np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz'), uv_mask = down)
+
+
+class UVMaskStore():
+    def __init__(self, folder, size):
+        self.size = size
+
+    def __getitem__(self, id) -> dict:
+        return uv_masks[self.size]
+
+
+def downscale_uv_maps_posed_body(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
+        img = fill(img, uv_masks_body[512])
+        for size in sizes:
+            mask = uv_masks_body[size]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)
+
+
+def downscale_uv_normals(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals']
+        img = fill(img, uv_masks_body[512])
+        for size in sizes:
+            mask = uv_masks_body[size]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_normals = down)
+
+
+
+def downscale_uv_maps_unposed(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
+        mask_full = uv_mask_store[512][None]
+        img = fill(img, mask_full)
+        for size in sizes:
+            mask = uv_mask_store[size][None]
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
+
+
+def downscale_uv_maps_static(folder, type, sizes):
+    for size in sizes:
+        os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
+        mask_full = uv_mask_store[512][None]
+        img = fill(img, mask_full)
+
+        for size in sizes:
+            down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy() 
+            mask = uv_mask_store[size][None]
+            down = down * mask[None,:,:]
+            np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)
+
+
+def downscale_uv_maps_mask(folder, type, size):
+    os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
+
+    frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
+        down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+        down = np.squeeze(down.numpy().astype(bool))
+        np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male')
+    parser.add_argument("-s", "--start", help = "start of range", default=0)
+    parser.add_argument("-e", "--end", help = "end of range", default=100000)
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
+    args = parser.parse_args()
+
+    device = 'cuda'
+
+    sizes = [32, 64, 128, 256]
+    folder = args.folder
+
+    resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
+    resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
+
+    img_idxs, img_coords, uv_masks, uv_masks_body = get_img_idxs()
+
+    uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
+    uv_mask_store[512] = UVMaskStore(folder, size=512)
+
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
+    # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
+    # downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
+
+    downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
+    downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
+    downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
+
+
+    print()

smpl/smpl_torch.py

@@ -10,7 +10,7 @@ from time import time
 from torch.nn import Module
 
 class SMPL(Module):
-  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
+  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False, batch_size=4):
     
     super(SMPL, self).__init__()
     self.J_regressor = []
@@ -56,6 +56,14 @@ class SMPL(Module):
     
     self.shape = (1,27578,3)
     self.device = device
+    self.batch_size = batch_size
+
+    self.with_zeros_ones = torch.tensor(
+        [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+      ).expand(self.batch_size,-1,-1).to(self.device)
+    self.results_zeros = torch.zeros((self.batch_size, 24, 1), dtype=torch.float32).to(self.device)
+    self.J_ones = torch.ones((self.batch_size, 24, 1), dtype=torch.float32, device=self.device)
+    self.rest_shape_h_ones_full = torch.ones((self.batch_size, self.weights_all.shape[1], 1), dtype=torch.float32).to(self.device)
 
     if orig_smpl:
         self.to_keep = self.to_keep[self.to_keep < 6890]
@@ -67,6 +75,8 @@ class SMPL(Module):
         self.weights = self.weights_all[:,self.to_keep,:]
         self.shape = (1, 6890, 3)
 
+    self.rest_shape_h_ones = torch.ones((self.batch_size, self.to_keep.shape[0], 1), dtype=torch.float32).to(self.device)
+
   @staticmethod
   def rodrigues(r):
     """
@@ -91,7 +101,7 @@ class SMPL(Module):
     m = torch.stack(
       (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
        -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
-    m = torch.reshape(m, (-1, 3, 3))
+    m = m.view((-1, 3, 3))
     i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
              + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
     A = r_hat.permute(0, 2, 1)
@@ -112,8 +122,8 @@ class SMPL(Module):
     Tensor after appending of shape [4,4]
     """
     ones = torch.tensor(
-      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+        [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
-    ).expand(x.shape[0],-1,-1).to(x.device)
+      ).expand(x.shape[0],-1,-1).to(x.device)
     ret = torch.cat((x, ones), dim=1)
     return ret
 
@@ -158,8 +168,8 @@ class SMPL(Module):
 
     # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
     v_shaped = []
-    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + self.v_template[0].view(self.shape)) # female
-    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + self.v_template[1].view(self.shape)) # male
     v_shaped = torch.stack(v_shaped, axis=-1)
     
     # J = torch.matmul(self.J_regressor, v_shaped)
@@ -167,7 +177,7 @@ class SMPL(Module):
       J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
     else:
       J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
-      J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+      J0 = J[:,0,:].view((-1, 1, 3))
       J = J - J0
     R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
 
@@ -177,12 +187,12 @@ class SMPL(Module):
       R_cube = R_cube_big[:, 1:, :, :]
       I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
         torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
-      lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
+      lrotmin = (R_cube - I_cube).view(batch_num, -1, 1).squeeze(dim=2)
       v_posed = torch.stack([v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[torch.greater(g, 0.5).to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
 
     results = []
     results.append(
-      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+      self.with_zeros(torch.cat((R_cube_big[:, 0], J[:, 0, :].view((-1, 3, 1))), dim=2))
     )
     for i in range(1, 24):
       results.append(
@@ -190,7 +200,7 @@ class SMPL(Module):
           results[self.parent[i]],
           self.with_zeros(
             torch.cat(
-              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              (R_cube_big[:, i], (J[:, i, :] - J[:, self.parent[i], :]).view((-1, 3, 1))),
               dim=2
             )
           )
@@ -202,25 +212,25 @@ class SMPL(Module):
       self.pack(
         torch.matmul(
           stacked,
-          torch.reshape(
+            torch.cat((J, self.results_zeros), dim=2).view(
-            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
             (batch_num, 24, 4, 1)
           )
         )
       )
     
-    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, self.J_ones], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
     weights = self.weights if verts is not None else self.weights_all 
+    rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full
     T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) 
     rest_shape_h = torch.cat(
-      (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
+      (v_posed, rest_shape_h_ones_local), dim=2
     )
-    v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
+    v = torch.matmul(T, rest_shape_h.view((batch_num, -1, 4, 1)))
-    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    v = v.view((batch_num, -1, 4))[:, :, :3]
     result = v
 
-    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    root_joint = joints[:,0,:].view((-1, 1, 3))
     result = result - root_joint
     joints = joints - root_joint