downscale multiple resolutions at the same time

This commit is contained in:
2022-10-02 13:35:13 +00:00
parent a80a8a868c
commit 52c139bbdd
+44 -29
View File
@@ -1,3 +1,4 @@
import argparse
import os import os
import tqdm import tqdm
import cv2 import cv2
@@ -7,7 +8,7 @@ import numpy as np
import plotly.graph_objects as go import plotly.graph_objects as go
import plotly.express as ex import plotly.express as ex
os.environ["CUDA_VISIBLE_DEVICES"] = "4" os.environ["CUDA_VISIBLE_DEVICES"] = "1"
gpus = tf.config.experimental.list_physical_devices('GPU') gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus: if gpus:
@@ -67,67 +68,76 @@ class UVMaskStore():
return self.store[id] return self.store[id]
def downscale_uv_maps_posed_body(folder, type, size: int): 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) os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type)) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames): for frame in tqdm.tqdm(frames):
id, _ = frame.split('_') id, _ = frame.split('_')
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient'] img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
mask_full = uv_mask_store_full[id] mask_full = uv_mask_store[512][id]
img = fill(img, mask_full) img = fill(img, mask_full)
mask = uv_mask_store_low[id] 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 = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
down = down * mask[:,:,np.newaxis] down = down * mask[:,:,np.newaxis]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down) 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): def downscale_uv_maps_curv(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type)) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames): for frame in tqdm.tqdm(frames):
id, _ = frame.split('_') id, _ = frame.split('_')
img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv'] img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
mask_full = uv_mask_store_full[id] mask_full = uv_mask_store[512][id]
img = fill_label(img, mask_full) img = fill_label(img, mask_full)
mask = uv_mask_store_low[id] 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 = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
down = np.squeeze(down) * mask[:,:] down = np.squeeze(down) * mask[:,:]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = down) np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_curvature = down)
def downscale_uv_maps_unposed(folder, type, size): def downscale_uv_maps_unposed(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type)) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames): for frame in tqdm.tqdm(frames):
id, _ = frame.split('_') id, _ = frame.split('_')
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape'] img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
mask_full = uv_mask_store_full[id] mask_full = uv_mask_store[512][id]
img = fill(img, mask_full) img = fill(img, mask_full)
mask = uv_mask_store_low[id] 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 = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
down = down * mask[:,:,np.newaxis] down = down * mask[:,:,np.newaxis]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down) np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
def downscale_uv_maps_static(folder, type, size): def downscale_uv_maps_static(folder, type, sizes):
for size in sizes:
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type)) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames): for frame in tqdm.tqdm(frames):
id, _ = frame.split('.') id, _ = frame.split('.')
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static'] img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
mask_full = uv_mask_store_full[id] mask_full = uv_mask_store[512][id]
img = fill(img, mask_full) img = fill(img, mask_full)
for size in sizes:
down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5) down = tf.image.resize(img, size=(size, int(size/2)), method=tf.image.ResizeMethod.LANCZOS5)
mask = uv_mask_store_low[id] mask = uv_mask_store[size][id]
down = down * mask[:,:,np.newaxis] down = down * mask[:,:,np.newaxis]
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down) np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)
@@ -135,7 +145,7 @@ def downscale_uv_maps_static(folder, type, size):
def downscale_uv_maps_mask(folder, type, size): def downscale_uv_maps_mask(folder, type, size):
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True) os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
frames = os.listdir(os.path.join(folder, type)) frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
for frame in tqdm.tqdm(frames): for frame in tqdm.tqdm(frames):
id, _ = frame.split('.') id, _ = frame.split('.')
@@ -147,20 +157,25 @@ def downscale_uv_maps_mask(folder, type, size):
if __name__ == '__main__': if __name__ == '__main__':
size = 32 parser = argparse.ArgumentParser()
folder = '/home/workspace/local-data/train' parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=1)
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
args = parser.parse_args()
uv_mask_store_full = UVMaskStore(folder, size=512)
uv_mask_store_low = UVMaskStore(folder, size=size) sizes = [32, 64, 128, 256]
folder = args.folder
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=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_static(folder=folder, type='uv_static', sizes=sizes)
# downscale_uv_maps_curv(folder=folder, type='uv_curvature', size=size) # downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
# downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', size=size) 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)
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() print()