208 lines
8.3 KiB
Python
208 lines
8.3 KiB
Python
import argparse
|
|
import os
|
|
import tqdm
|
|
import cv2
|
|
import tensorflow as tf
|
|
import scipy.interpolate as spi
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
|
|
|
gpus = tf.config.experimental.list_physical_devices('GPU')
|
|
if gpus:
|
|
try:
|
|
# Currently, memory growth needs to be the same across GPUs
|
|
for gpu in gpus:
|
|
tf.config.experimental.set_memory_growth(gpu, True)
|
|
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
|
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
|
except RuntimeError as e:
|
|
# Memory growth must be set before GPUs have been initialized
|
|
print(e)
|
|
|
|
def get_img_idxs():
|
|
img_idxs = dict()
|
|
img_coords = dict()
|
|
uv_masks = dict()
|
|
uv_mask = np.load('custom_utils/mask_256x512.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 = tf.image.resize(uv_mask[:,:,tf.newaxis].astype(np.int32), size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
|
down = np.squeeze(down.numpy().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)
|
|
|
|
return img_idxs, img_coords, uv_masks
|
|
|
|
|
|
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):
|
|
id, _ = frame.split('_')
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
|
|
mask_full = uv_mask_store[512][id]
|
|
img = fill(img, mask_full)
|
|
for size in sizes:
|
|
mask = uv_mask_store[size][id]
|
|
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
|
down = down * mask[:,:,np.newaxis]
|
|
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)
|
|
|
|
|
|
def downscale_uv_maps_curv(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):
|
|
id, _ = frame.split('_')
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
|
|
mask_full = uv_mask_store[512][id]
|
|
img = fill_label(img, mask_full)
|
|
for size in sizes:
|
|
mask = uv_mask_store[size][id]
|
|
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
|
down = np.squeeze(down) * mask[:,:]
|
|
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = 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):
|
|
id, _ = frame.split('_')
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
|
|
mask_full = uv_mask_store[512][id]
|
|
img = fill(img, mask_full)
|
|
for size in sizes:
|
|
mask = uv_mask_store[size][id]
|
|
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
|
down = down * mask[:,:,np.newaxis]
|
|
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):
|
|
id, _ = frame.split('.')
|
|
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
|
|
mask_full = uv_mask_store[512][id]
|
|
img = fill(img, mask_full)
|
|
|
|
for size in sizes:
|
|
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
|
|
mask = uv_mask_store[size][id]
|
|
down = down * mask[:,:,np.newaxis]
|
|
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):
|
|
id, _ = frame.split('.')
|
|
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='/home/workspace/local-data/cape_processed/test/male')
|
|
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
|
parser.add_argument("-e", "--end", help = "end of range", default=10000)
|
|
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
|
args = parser.parse_args()
|
|
|
|
|
|
sizes = [32, 64, 128, 256]
|
|
folder = args.folder
|
|
|
|
img_idxs, img_coords, uv_masks = 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_maps_curv(folder=folder, type='uv_curvature', sizes=sizes)
|
|
# downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
|
|
|
|
|
|
print() |