mask broadcasting

eduardorittner · Jun 29, 2024 · ef1e7ec · ef1e7ec
1 parent f69cb43
commit ef1e7ec
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/run_dwt.py b/run_dwt.py
@@ -134,7 +134,21 @@ def run_dwt(
                 )
                 exit(1)
 
+            mask_id = mask_files[i].split("-")[0]
+
+            if mask_id != id:
+                print("[ERROR]: Mask id ({mask_id}) and image id ({id}) must be the same.")
+                exit(1)
+
+            print(mask)
+
+            mask = np.expand_dims(mask, np.argmin(result.shape))
+
+            print(mask.shape)
+
+            print(result)
             result = result * mask
+            print(result)
             expected = expected * mask
 
         if "ssim" in metrics: