Opt: Mask.apply on both monochrome and RGB images

- Opt: Cached mask file in sensitive info hiding

module/base/mask.py

@@ -5,6 +5,17 @@ from PIL import Image
 from module.base.template import Template
 
 
+def get_image_channel(image):
+    """
+    Args:
+        image (np.ndarray):
+
+    Returns:
+        int: 0 for monochrome, 3 for RGB.
+    """
+    return 3 if len(image.shape) == 3 else 0
+
+
 class Mask(Template):
     @property
     def image(self):
@@ -13,5 +24,42 @@ class Mask(Template):
 
         return self._image
 
+    @image.setter
+    def image(self, value):
+        self._image = value
+
+    def set_channel(self, channel):
+        """
+        Args:
+            channel (int): 0 for monochrome, 3 for RGB.
+
+        Returns:
+            bool: If changed.
+        """
+        image_channel = get_image_channel(self.image)
+        if channel == 0:
+            if image_channel == 0:
+                return False
+            else:
+                self._image, _, _ = cv2.split(self._image)
+                return True
+        else:
+            if image_channel == 0:
+                self._image = cv2.merge([self._image] * 3)
+                return True
+            else:
+                return False
+
     def apply(self, image):
+        """
+        Apply mask on image.
+
+        Args:
+            image:
+
+        Returns:
+            np.ndarray:
+        """
+        image = np.array(image)
+        self.set_channel(get_image_channel(image))
         return cv2.bitwise_and(image, self.image)

module/handler/sensitive_info.py

@@ -2,8 +2,12 @@ import re
 
 from PIL import Image
 
+from module.base.mask import Mask
 from module.ui.page import *
 
+MASK_MAIN = Mask('./assets/mask/MASK_MAIN.png')
+MASK_PLAYER = Mask('./assets/mask/MASK_PLAYER.png')
+
 
 def put_image_mask(image, mask):
     """
@@ -29,11 +33,11 @@ def handle_sensitive_image(image):
         PIL.Image.Image:
     """
     if PLAYER_CHECK.match(image, offset=(30, 30)):
-        return put_image_mask(image, mask='MASK_PLAYER')
+        image = MASK_PLAYER.apply(image)
     if MAIN_CHECK.match(image, offset=(30, 30)):
-        return put_image_mask(image, mask='MASK_MAIN')
+        image = MASK_MAIN.apply(image)
 
-    return image
+    return Image.fromarray(image, mode='RGB')
 
 
 def handle_sensitive_text(text):