From 97d53fd2f3df0e6bf1fbd43bf44178ff4a1fac07 Mon Sep 17 00:00:00 2001
From: cat <cat@wolfgirl.org>
Date: Sun, 4 Nov 2018 20:24:16 +0500
Subject: [PATCH] Optimize mask computation in bar decensoring mode.

Use numpy vector operations for massive speedup, ~30x on 1000x1500 image.
---
 decensor.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/decensor.py b/decensor.py
index 6bd4d0f..2650096 100644
--- a/decensor.py
+++ b/decensor.py
@@ -21,14 +21,10 @@ class Decensor:
 
         self.load_model()
 
-    def get_mask(self, colored, width, height):
+    def get_mask(self, colored):
         mask = np.ones(colored.shape, np.uint8)
-        #count = 0
-        #TODO: change to iterate over all images in batch when implementing batches
-        for row in range(height):
-            for col in range(width):
-                if np.array_equal(colored[0][row][col], self.mask_color):
-                    mask[0, row, col] = 0
+        i, j = np.where(np.all(colored[0] == self.mask_color, axis=-1))
+        mask[0, i, j] = 0
         return mask
 
     def load_model(self):
@@ -93,7 +89,7 @@ class Decensor:
             #if mosaic decensor, mask is empty
             mask = np.ones(ori_array.shape, np.uint8)
         else:
-            mask = self.get_mask(ori_array, width, height) 
+            mask = self.get_mask(ori_array)
 
         #colored image is only used for finding the regions
         regions = find_regions(colored.convert('RGB'))