only read masks once for downscaling

downscale.py

@@ -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']
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
-        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
+        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']
+        img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv']
-        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
+        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']
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
-        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
+        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']
+        img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
-        mask = np.load(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz')); mask = mask['uv_mask']
+        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)
+    folder = '/home/workspace/local-data/train'
-    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)
 
+    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()