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"] = "4" 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 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 class UVMaskStore(): def __init__(self, folder, size): self.store = dict() self.folder = folder self.size = size def read(self, id): mask = np.load(os.path.join(self.folder, 'uv_mask' + ('_' + str(self.size) if self.size != 512 else ''), id + '.npz')) mask = mask['uv_mask'] return mask def __getitem__(self, id) -> dict: if id not in self.store: self.store[id] = self.read(id) return self.store[id] def downscale_uv_maps_posed_body(folder, type, size: int): os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type)) 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_full[id] img = fill(img, mask_full) mask = uv_mask_store_low[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, size): os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type)) 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_full[id] img = fill_label(img, mask_full) mask = uv_mask_store_low[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, size): os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type)) 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_full[id] img = fill(img, mask_full) mask = uv_mask_store_low[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, size): os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) frames = os.listdir(os.path.join(folder, type)) 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_full[id] img = fill(img, mask_full) down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5) mask = uv_mask_store_low[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)) 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__': size = 32 folder = '/home/workspace/local-data/train' uv_mask_store_full = UVMaskStore(folder, size=512) uv_mask_store_low = UVMaskStore(folder, size=size) # downscale_uv_maps_mask(folder=folder, type='uv_mask', size=size) # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', size=size) # downscale_uv_maps_curv(folder=folder, type='uv_curvature', size=size) # downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', size=size) downscale_uv_maps_static(folder=folder, type='uv_static', size=32) downscale_uv_maps_static(folder=folder, type='uv_static', size=64) downscale_uv_maps_static(folder=folder, type='uv_static', size=128) downscale_uv_maps_static(folder=folder, type='uv_static', size=256) print()