103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
import os
|
|
import tqdm
|
|
import scipy.interpolate as spi
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
|
|
|
|
def fill(img, mask):
|
|
coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
|
img_idx = np.transpose(np.stack(np.nonzero((mask))), (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 = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
|
|
img_idx = np.transpose(np.stack(np.nonzero((mask))), (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), dtype=np.float32)
|
|
res[coords[:,0], coords[:,1]] = img_2
|
|
res[img_idx[:,0], img_idx[:,1]] = img_data
|
|
|
|
return res
|
|
|
|
|
|
def fill_uv_maps_posed_body(folder, type):
|
|
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
|
mask = np.load('custom_utils/mask_256x512.npy')
|
|
|
|
frames = os.listdir(os.path.join(folder, type))
|
|
|
|
for frame in tqdm.tqdm(frames):
|
|
id, _ = frame.split('_')
|
|
# mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
|
|
filled = fill(img, mask)
|
|
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)
|
|
|
|
|
|
def fill_uv_maps_curv(folder, type):
|
|
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
|
mask = np.load('custom_utils/mask_256x512.npy')
|
|
|
|
frames = os.listdir(os.path.join(folder, type))
|
|
|
|
for frame in tqdm.tqdm(frames):
|
|
id, _ = frame.split('_')
|
|
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
|
|
filled = fill_label(img, mask)
|
|
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_curvature_filled = filled)
|
|
|
|
|
|
|
|
def fill_uv_maps_unposed(folder, type):
|
|
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
|
mask = np.load('custom_utils/mask_256x512.npy')
|
|
|
|
frames = os.listdir(os.path.join(folder, type))
|
|
|
|
for frame in tqdm.tqdm(frames):
|
|
id, _ = frame.split('_')
|
|
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
|
|
filled = fill(img, mask)
|
|
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_unposed_no_shape_filled = filled)
|
|
|
|
|
|
|
|
def fill_uv_maps_static(folder, type):
|
|
os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True)
|
|
mask = np.load('custom_utils/mask_256x512.npy')
|
|
|
|
frames = os.listdir(os.path.join(folder, type))
|
|
|
|
for frame in tqdm.tqdm(frames):
|
|
id, _ = frame.split('.')
|
|
mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
|
|
filled = fill(img, mask)
|
|
np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_static_filled = filled)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# fill_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static')
|
|
# fill_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape')
|
|
# fill_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature')
|
|
fill_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient')
|
|
|
|
|
|
print() |