luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

only read masks once for downscaling

Browse Source
This commit is contained in:
luni
2022-09-30 21:50:32 +00:00
parent e3e7087c10
commit 6a2f9ff4ff
1 changed files with 59 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
downscale.py
+59 -14
View File
@@ -7,8 +7,19 @@ import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
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))
@@ -38,6 +49,24 @@ def fill_label(img, mask):
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)
@@ -45,8 +74,10 @@ def downscale_uv_maps_posed_body(folder, type, size: int):
for frame in tqdm.tqdm(frames):
id, _ = frame.split('_')
img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_posed_body_no_orient_filled']
mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
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)
@@ -59,8 +90,10 @@ def downscale_uv_maps_curv(folder, type, size):
for frame in tqdm.tqdm(frames):
id, _ = frame.split('_')
img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_curvature_filled']
mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
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)
@@ -74,8 +107,10 @@ def downscale_uv_maps_unposed(folder, type, size):
for frame in tqdm.tqdm(frames):
id, _ = frame.split('_')
img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_unposed_no_shape_filled']
mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
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)
@@ -88,9 +123,11 @@ def downscale_uv_maps_static(folder, type, size):
for frame in tqdm.tqdm(frames):
id, _ = frame.split('.')
img = np.load(os.path.join(folder, type + '_filled', frame)); img = img['uv_static_filled']
mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
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)
@@ -111,11 +148,19 @@ def downscale_uv_maps_mask(folder, type, size):
if __name__ == '__main__':
size = 32
# downscale_uv_maps_mask(folder='/home/workspace/local-data/train_100', type='uv_mask', size=size)
downscale_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static', size=size)
downscale_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape', size=size)
downscale_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature', size=size)
downscale_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient', size=size)
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()