diff --git a/downscale.py b/downscale.py index 411615f..31ac1b9 100644 --- a/downscale.py +++ b/downscale.py @@ -8,7 +8,7 @@ import numpy as np import plotly.graph_objects as go import plotly.express as ex -os.environ["CUDA_VISIBLE_DEVICES"] = "1" +os.environ["CUDA_VISIBLE_DEVICES"] = "3" gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: @@ -50,6 +50,13 @@ def fill_label(img, mask): 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.store = dict() @@ -57,7 +64,11 @@ class UVMaskStore(): 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')) + path = os.path.join(self.folder, 'uv_mask' + ('_' + str(self.size) if self.size != 512 else ''), id + '.npz') + if not os.path.exists(path): + generate_downscaled_uv_map(self.folder, id, self.size) + + mask = np.load(path) mask = mask['uv_mask'] return mask @@ -158,9 +169,9 @@ def downscale_uv_maps_mask(folder, type, size): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt') + parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/local-data/valid') 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("-e", "--end", help = "end of range", default=1000) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') args = parser.parse_args() @@ -171,7 +182,10 @@ if __name__ == '__main__': 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=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)