fix static cache

data/convert_to_uv_cape.py

@@ -0,0 +1,530 @@
+import os
+import sys
+import glob
+import tqdm
+import math
+import time
+import random
+import shutil
+import collections
+import argparse
+import pickle
+import numpy as np
+import scipy as sp
+import trimesh as tm
+# import cmapy as cm
+import tables as tb
+import scipy.io as sio
+import scipy.interpolate as spi
+import matplotlib.pyplot as plt
+import scipy.stats as stat
+import cv2
+import plotly.express as ex
+import plotly.graph_objects as go
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures
+from data.cloth3d_io import readPC2
+
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+
+def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None):
+    verts_uv = np.zeros((res * 2, res, 3))
+    uv_values = np.zeros((uv_map.shape[0], 3))
+    uv_values[m2uv[:, 0]] = verts
+    uv_values[m2uv[:, 1]] += verts
+    uv_values = uv_values[uv_map_mask[type], :] #   add garment mask
+    uv_values /= uv_counter[type][:, np.newaxis]
+    return_garment_uv_mask = False
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask[type]] #   add garment mask
+    else:
+        garment_mask = np.sum(uv_values, axis=-1) != 0
+
+    # test
+    verts_uv[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
+    # verts_uv_save = verts_uv.copy()
+
+    # garment_mask = np.sum(uv_values, axis=-1) != 0
+    unique_uv_values = uv_values[garment_mask, :]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
+    rbf_dest_current = rbf_dest[type]
+
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4)
+        garment_uv_mask_border = garment_uv_mask.astype(bool)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
+        garment_uv_mask = garment_uv_mask * uv_mask[type][:,:,0]
+        rbf_dest_current = rbf_dest[type] * garment_uv_mask
+        return_garment_uv_mask = True
+
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
+    if method == 'rbf':
+        interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear')
+    elif method == 'nearest':
+        interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values)
+    else:
+        raise NotImplementedError
+
+    final = interpolator(rbf_dest_current)
+
+    result = np.zeros((res * 2, res, 3))
+
+    # not sure, shouldn't be needed
+    # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
+
+    result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
+    result[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
+    if garment_uv_mask is not None:
+        result = result * garment_uv_mask[:,:,np.newaxis]
+
+    if return_garment_uv_mask:
+        return result, garment_uv_mask, garment_uv_mask_border
+    else:
+        return result 
+
+
+def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
+    batch = verts.shape[0]
+    verts_uv = np.zeros((batch, res * 2, res, 3))
+    uv_values = np.zeros((batch, uv_map.shape[0], 3))
+    uv_values[:, m2uv[:, 0]] = verts
+    uv_values[:, m2uv[:, 1]] += verts
+    uv_values = uv_values[:, uv_map_mask['non-skirt'], :]
+    uv_values /= uv_counter['non-skirt'][np.newaxis, :, np.newaxis]
+    return_garment_uv_mask = False
+    # test
+    verts_uv[:, uv_map_discreet_unique['non-skirt'][:, 0], uv_map_discreet_unique['non-skirt'][:, 1], :] = uv_values
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask['non-skirt']] #   add garment mask
+    else:
+        garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
+
+    unique_uv_values = uv_values[:, garment_mask, :]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique['non-skirt'][garment_mask, :]
+    rbf_dest_current = rbf_dest['non-skirt']
+
+    # add back mesh vertices that are not part of the garment
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique['non-skirt'][np.invert(garment_mask), :]
+    rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
+    rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
+    
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        return_garment_uv_mask = True
+
+    if garment_uv_mask is not None:
+        rbf_dest_current = (rbf_dest['non-skirt'] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    final = []
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+        if method == 'rbf':
+            interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear')
+        elif method == 'nearest':
+            interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:])
+        else:
+            raise NotImplementedError
+
+        result = np.zeros((res * 2, res, 3), np.float32)
+        result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
+        result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[i,:,:]
+
+        if garment_uv_mask is not None:
+            result *= garment_uv_mask[:,:,np.newaxis]
+        
+        final.append(result)
+
+    if return_garment_uv_mask:
+        return final, garment_uv_mask
+    else:
+        return final 
+    
+
+def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None, id=None):
+    batch = len(verts)
+
+    assert method == 'rbf'
+
+    assert mask is not None
+    garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+    garment_mask[m2uv[:, 0]] = mask
+    garment_mask[m2uv[:, 1]] += mask
+    garment_mask = garment_mask[uv_map_mask[type]] #   add garment mask
+
+    rbf_dest_current = rbf_dest[type]
+
+    # add back mesh vertices that are not part of the garment
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[type][np.invert(garment_mask), :]
+    rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
+    rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
+
+    assert not (garment_uv_mask is None and mask is not None)
+    assert garment_uv_mask is not None
+    rbf_dest_current = (rbf_dest[type] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    out = []
+
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+
+        uv_values = np.zeros((uv_map.shape[0], 3))
+        uv_values[m2uv[:, 0]] = verts[i]
+        uv_values[m2uv[:, 1]] += verts[i]
+        uv_values = uv_values[uv_map_mask[type], :]
+        uv_values /= uv_counter[type][:, np.newaxis]
+
+        unique_uv_values = uv_values[garment_mask, :]
+        uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
+
+        interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
+        result = np.zeros((res * 2, res, 3), np.float32)
+        result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = interpolator(rbf_dest_current)
+        result[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values[:,:]
+
+        if garment_uv_mask is not None:
+            result *= garment_uv_mask[:,:,np.newaxis]
+        
+        out.append(result)
+    
+    return out
+
+
+def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None):
+    batch = values.shape[0]
+    values = (values + 1) * mask
+    values_uv = np.zeros((batch, res * 2, res))
+    missing = np.zeros((res * 2, res))
+    uv_values = np.zeros((batch, uv_map.shape[0]))
+    uv_values[:, m2uv[:, 0]] = values
+    uv_values[:, m2uv[:, 1]] = values
+    uv_values = uv_values[:, uv_map_mask]
+    return_garment_uv_mask = False
+
+    if mask is not None:
+        garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
+        garment_mask[m2uv[:, 0]] = mask
+        garment_mask[m2uv[:, 1]] += mask
+        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    else:
+        raise Exception
+
+
+    # test 
+    values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
+    missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0)
+    
+    unique_uv_values = uv_values[:, garment_mask]
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+    # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0]
+
+    if garment_uv_mask is None and mask is not None:
+        garment_uv_mask = np.zeros((res * 2, res))
+        garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255
+        garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4)
+        garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
+        return_garment_uv_mask = True
+
+    rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask
+    rbf_dest_current = np.nonzero(rbf_dest_current)
+    rbf_dest_current = np.array(rbf_dest_current).T
+
+    # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T
+    final = []
+    for i in range(batch):
+        print('processing frame: ' + str(i))
+        interpolator =  spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:])
+        final.append(interpolator(rbf_dest_current))
+    final = np.asarray(final)
+
+    result = np.zeros((batch, res * 2, res))
+    result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values
+    result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final
+    if garment_uv_mask is not None:
+        result = result * garment_uv_mask[np.newaxis,:,:]
+
+    if return_garment_uv_mask:
+        return result, garment_uv_mask
+    else:
+        return result 
+
+
+def m2uv_array(m2uv):
+    if m2uv is None:
+        return None
+    result = np.zeros((3638, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv, m2uv):
+    res = uv.shape[0]
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+    pass
+
+
+def verify_reconstruction(orig, reconstructed):
+    error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1))
+    colors = np.array([cm.color('jet', 1 - i) for i in error])
+    error_sum = np.sum(error)
+
+
+    verts_uv = np.zeros((uv_res * 2, uv_res, 3))
+    uv_values = np.zeros((uv_map.shape[0], 3))
+    uv_values[m2uv[:, 0]] = colors
+    uv_values[m2uv[:, 1]] += colors
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_values /= uv_counter[:, np.newaxis]
+    verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(111)
+    ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255)
+    plt.show()
+
+
+
+    fig = plt.figure(figsize=(15, 15))
+    ax = fig.add_subplot(121, projection='3d')
+    x = orig
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5)
+
+    ax = fig.add_subplot(122, projection='3d')
+    x = reconstructed#[error > 0.1, :]
+    ax.set_xlim3d(-1, 1)
+    ax.set_ylim3d(-1, 1)
+    ax.set_zlim3d(-1, 1)
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :]
+
+    plt.show()
+
+    pass
+
+# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs
+def unique_pairs(a):
+    a.sort(axis=1)
+    b = np.ascontiguousarray(a).view(
+        np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
+    )
+    _, inv, ct = np.unique(b, return_inverse=True, return_counts=True)
+    return ct[inv] == 1
+
+
+def unique_uv_map_mask(uv_map_discreet):
+    mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
+    temp = np.zeros((uv_res * 2, uv_res))
+    for p in uv_map_discreet:
+        temp[p[0], p[1]] += 1
+    # temp = temp == 1
+    for i, p in enumerate(uv_map_discreet):
+        if temp[p[0], p[1]] == 1:
+            mask[i] = True
+        elif temp[p[0], p[1]] == 2:
+            mask[i] = True
+            temp[p[0], p[1]] = 0
+    return mask
+
+
+def test(verts, faces, values):
+    uv_values = np.zeros((uv_map.shape[0], 3))
+    uv_values[m2uv[:, 0]] = values
+    uv_values[m2uv[:, 1]] += values
+    uv_counter = np.zeros((uv_map.shape[0]))
+    uv_counter[m2uv[:, 0]] = 1
+    uv_counter[m2uv[:, 1]] += 1
+    uv_values /= uv_counter[:, np.newaxis]
+    x = uv_values
+    verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
+    im = ((x - x.min()) / (x.max() - x.min())) * 255
+    tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show()
+
+
+def get_shapedirs():
+    to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1)
+    shapedirs = []
+    with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        shapedirs.append(female['shapedirs'][to_keep, :])
+    with open(os.path.join('smpl/super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        shapedirs.append(male['shapedirs'][to_keep, :])
+    return shapedirs
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/test/female')
+    parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
+    parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
+    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()
+
+    F = {'non-skirt': np.load("smpl/smpl_faces_orig.npy"), 'skirt': None}
+    uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3_orig.npy'), 'skirt': None}
+    m2uv = {'non-skirt': np.load('custom_utils/m2uv_orig.npy', allow_pickle=True).item(), 'skirt': None}
+    m2uv = {k: m2uv_array(v) for k, v in m2uv.items()}
+    uv_mask = {'non-skirt': np.load('custom_utils/mask_256x512.npy'), 'skirt': np.load('custom_utils/mask_256x512_skirt.npy')}
+
+
+    uv_res = 256
+    uv_map_discreet = dict()
+    uv_map_discreet_unique = dict()
+    uv_counter = dict()
+    rbf_dest = dict()
+    uv_map_mask = dict()
+    for k,v in uv_mask.items():
+        if k == 'skirt':
+            continue
+        uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+        # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
+        # uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
+        uv_mask[k] = cv2.dilate(v.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+
+        uv_map_discreet[k] = np.floor(uv_map[k] * uv_res).astype(np.int32)
+        uv_map_mask[k] = unique_uv_map_mask(uv_map_discreet[k])
+
+        uv_counter[k] = np.zeros((uv_map[k].shape[0]))
+        uv_counter[k][m2uv[k][:, 0]] = 1
+        uv_counter[k][m2uv[k][:, 1]] += 1
+        uv_counter[k] = uv_counter[k][uv_map_mask[k]]
+
+        uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :]
+
+        rbf_dest[k] = np.ones((uv_res * 2, uv_res))
+        rbf_dest[k][uv_map_discreet_unique[k][:, 0], uv_map_discreet_unique[k][:, 1]] = 0
+        rbf_dest[k][np.invert(uv_mask[k][:,:,0])] = 0
+
+
+    #########################################################################################
+
+    folder = args.folder
+    # data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
+    data_to_convert = ['mesh_static', 'mesh_normals'] # 'mesh_unposed',
+
+    frames = sorted(os.listdir(folder + '/mesh_posed_body_no_orient'))#[args.start:args.end]
+    frames = [[frame.split('.')[0], frame.split('.')[1]] for frame in frames]
+    ids = list(set([frame[0] for frame in frames]))
+    frames_by_id = {id: [frame[1] for frame in frames if frame[0] == id] for id in ids}
+
+    shapedirs = get_shapedirs()
+
+    root = '/'.join(folder.split('/')[:-1])
+    type = 'non-skirt'
+
+    os.makedirs(folder, exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
+
+    mask = np.ones((3638), dtype=bool)
+
+    for id in tqdm.tqdm(ids):
+
+
+        # x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+        # mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+        # mesh_posed[:,mask, :] = x
+        if 'mesh_unposed' in data_to_convert:
+            mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+            mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
+        if 'mesh_static' in data_to_convert:
+            mesh_static = np.load(os.path.join(folder, 'mesh_static', f'{id}.npz'))['mesh_static']
+        if 'mesh_normals' in data_to_convert:
+            mesh_normals = np.stack([np.load(os.path.join(folder, 'mesh_normals', f'{id}.{int(i):06d}.npz'))['mesh_normals'] for i in frames_by_id[id]], axis=0)
+        if 'body_posed' in data_to_convert and k != 1:
+            body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
+
+        start = time.perf_counter()
+
+        uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv[type], uv_map[type], mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
+        uv_static = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        uv_static = uv_static[0].astype(np.float32)
+
+        if 'mesh_posed' in data_to_convert:
+            uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        if 'mesh_unposed' in data_to_convert:
+            uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        if 'mesh_normals' in data_to_convert:
+            uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+        if 'body_posed' in data_to_convert and k != 1:
+            uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
+
+        print(time.perf_counter() - start)
+
+        # ex.imshow(uv_mask_local).show()
+        # ex.imshow((uv_static+1)/2).show()
+        # ex.imshow((uv_posed[0]+1)/2).show()
+        # ex.imshow((uv_unposed[0]+1)/2).show()
+        # ex.imshow((uv_normals[0]+1.5)/3).show()
+        # ex.imshow((uv_body_posed[0]+1)/2).show()
+
+        # pass
+        # np.savez_compressed(os.path.join(folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
+        # np.savez_compressed(os.path.join(folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
+
+        if 'mesh_posed' in data_to_convert:
+            for i in range(len(uv_posed)):
+                np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient',  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
+        if 'mesh_unposed' in data_to_convert:
+            for i in range(len(uv_unposed)):
+                np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape',  f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
+        if 'mesh_normals' in data_to_convert:
+            for i, frame_nr in enumerate(frames_by_id[id]):
+                np.savez_compressed(os.path.join(folder, 'uv_normals',  f'{id}.{frame_nr}.npz'), uv_normals=uv_normals[i])
+        if 'body_posed' in data_to_convert and k != 1:
+            for i in range(len(uv_body_posed)):
+                np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
+        
+        # pass

data/dataset.py

@@ -105,15 +105,15 @@ class StaticCache():
 
 
     def __getitem__(self, frame) -> dict:
-        if frame not in self.store:
+        if (frame.id, frame.garment) not in self.store:
             item = self.read(frame.id, frame.garment)
             stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
             for key, value in stored_data.items():
                 if key.startswith('uv_') and value.shape[-1] == 3:
                     stored_data[key] = torch.permute(value, [2,0,1])
-            self.store[frame] = stored_data
+            self.store[(frame.id, frame.garment)] = stored_data
         
-        data = { key: value.to(self.device) for key, value in self.store[frame].items()}
+        data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()}
         return data
 
 

data/preprocessing.py

@@ -6,8 +6,8 @@ parser = argparse.ArgumentParser()
 # parser.add_argument("-f", "--folder", help = "end of range")
 parser.add_argument("-s", "--start", help = "start of range", type=int, default=0)
 parser.add_argument("-e", "--end", help = "end of range", type=int, default=1)
-parser.add_argument("-g", "--gpu", help = "gpu id")
-# parser.add_argument("-n", "--name", help = "end of range")
+parser.add_argument("-g", "--gpu", help = "gpu id", default='0')
+parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
 args = parser.parse_args()
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
@@ -356,6 +356,61 @@ def create_mesh_normals(folder, out_folder):
 
             pass
 
+def create_mesh_normals_cape(folder, out_folder):
+    frames = os.listdir(os.path.join(folder, 'mesh_posed'))
+    os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
+    F = np.load('smpl/smpl_faces_orig.npy')
+    # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    for frame in tqdm(frames):
+        id, frame_nr = frame.split('.')[:2]
+
+
+        verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}.{frame_nr}.npz'))['mesh_posed']
+        normals = calculate_normals(verts[None,:,:], F)[0]
+
+        np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}.{frame_nr}.npz'), mesh_normals=normals)
+
+        # normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
+        # colors = np.clip(np.abs(normals), 0, 1)
+        # verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
+        # plot as scatter plot with plotly
+        # import plotly.graph_objects as go
+        # frame = 5
+        # fig = go.Figure(data=[go.Scatter3d(
+        #     x=verts[:,0],
+        #     y=verts[:,2],
+        #     z=verts[:,1],
+        #     mode='markers',
+        #     marker=dict(
+        #         size=1,
+        #         color=normals[:,:],
+        #         colorscale='Viridis',
+        #         opacity=0.8
+        #     )
+        # )])
+        # moved_verts = verts + 0.05 * normals
+
+        # # Create the scatter plot for the moved vertices
+        # fig.add_trace(go.Scatter3d(
+        #     x=moved_verts[:,0],
+        #     y=moved_verts[:,2],
+        #     z=moved_verts[:,1],
+        #     mode='markers',
+        #     marker=dict(
+        #         size=1,
+        #         color='red',
+        #         opacity=0.8
+        #     )
+        # ))
+        # fig.show()
+
+        # # create a plotly mesh plot for verts and faces
+        # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+        # fig.show()
+
+        pass
+
 
 def verify_registration_ans_mask(folder, out_folder):
     ids = os.listdir(folder)
@@ -588,4 +643,7 @@ if __name__ == '__main__':
     # create_mesh_normals(folder, out_folder)
     # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
 
-    create_frames_file(folder, out_folder)
+    # create_frames_file(folder, out_folder)
+    
+    folder = '/data-local/127_2/cape_processed/test/female'
+    create_mesh_normals_cape(folder, None)