error proof downscaling

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)