Added support for custom mask colors

config.py

@@ -1,5 +1,14 @@
 import argparse
 
+def str2floatarr(v):
+    if type(v) == str:
+        try:
+            return [float(v) for v in v.split(',')]
+        except:
+            raise argparse.ArgumentTypeError('Integers seperated by commas expected.')
+    else:
+        raise argparse.ArgumentTypeError('Integers seperated by commas expected.')
+
 def str2bool(v):
     if v.lower() in ('yes', 'true', 't', 'y', '1', True):
         return True
@@ -17,10 +26,7 @@ def get_args():
 	parser.add_argument('--decensor_output_path', dest='decensor_output_path', default='./decensor_output/', help='output images generated from running decensor.py path')
 
 	#Decensor settings
-	parser.add_argument('--mask_color_red', dest='mask_color_red', default=0, help='red channel of mask color in decensoring')
-	parser.add_argument('--mask_color_green', dest='mask_color_green', default=255, help='green channel of mask color in decensoring')
-	parser.add_argument('--mask_color_blue', dest='mask_color_blue', default=0, help='blue channel of mask color in decensoring')
-
+	parser.add_argument('--mask_color', dest='mask_color', default=[0,255,0], type=str2floatarr, help='rgb color of the mask, comma seperated.')
 	parser.add_argument('--is_mosaic', dest='is_mosaic', default='False', type=str2bool, help='true if image has mosaic censoring, false otherwise')
 
 	args = parser.parse_args()

decensor.py

@@ -21,7 +21,7 @@ class Decensor:
         self.args = config.get_args()
         self.is_mosaic = self.args.is_mosaic
 
-        self.mask_color = [self.args.mask_color_red/255.0, self.args.mask_color_green/255.0, self.args.mask_color_blue/255.0]
+        self.mask_color = [float(v/255) for v in self.args.mask_color] # normalize mask color
 
         if not os.path.exists(self.args.decensor_output_path):
             os.makedirs(self.args.decensor_output_path)
@@ -71,7 +71,7 @@ class Decensor:
                             self.decensor_image(ori_img, colored_img, file_name)
                             break
                     else: #for...else, i.e if the loop finished without encountering break
-                        print("Corresponding original, uncolored image not found in {}.".format(ori_file_path))
+                        print("Corresponding original, uncolored image not found in {}.")
                         print("Check if it exists and is in the PNG or JPG format.")
                 else:
                     self.decensor_image(colored_img, colored_img, file_name)
@@ -111,7 +111,7 @@ class Decensor:
             mask = self.get_mask(ori_array)
 
         #colored image is only used for finding the regions
-        regions = find_regions(colored.convert('RGB'))
+        regions = find_regions(colored.convert('RGB'), [v*255 for v in self.mask_color]) #unnormalize the color so it can check against pixels
         print("Found {region_count} censored regions in this image!".format(region_count = len(regions)))
 
         if len(regions) == 0 and not self.is_mosaic:

docs/USAGE.md

@@ -2,6 +2,7 @@
 ## I. Decensoring bar censors
 
 For each image you want to decensor, using image editing software like Photoshop or GIMP to color the areas you want to decensor the green color (0,255,0), which is a very bright green color and be sure that the picture you use had RGB color Mode not Indexed Color Mode.
+Note: you can also specify a custom mask color, with the --mask_color option. For example, for bright red, you would use --mask_color=255,0,0
 
 *I strongly recommend you use the pencil tool and NOT the brush tool.*
 

libs/utils.py

@@ -3,17 +3,17 @@ from PIL import Image, ImageDraw
 
 #convert PIL image to numpy array
 def image_to_array(image):
-	array = np.asarray(image)
-	return np.array(array / 255.0)
+    array = np.asarray(image)
+    return np.array(array / 255.0)
 
 #find strongly connected components with the mask color
-def find_regions(image):
+def find_regions(image, mask_color):
     pixel = image.load()
     neighbors = dict()
     width, height = image.size
     for x in range(width):
         for y in range(height):
-            if is_green(pixel[x,y]):
+            if is_right_color(pixel[x,y], *mask_color):
                 neighbors[x, y] = {(x,y)}
     for x, y in neighbors:
         candidates = (x + 1, y), (x, y + 1)
@@ -122,14 +122,14 @@ def expand_bounding(img, region, expand_factor=1.5, min_size = 256):
         x1_square, y1_square, x2_square, y2_square = min_x, min_y, max_x, max_y
     return x1_square, y1_square, x2_square, y2_square
 
-def is_green(pixel):
-    r, g, b = pixel
-    return r == 0 and g == 255 and b == 0
+def is_right_color(pixel, r2, g2, b2):
+    r1, g1, b1 = pixel
+    return r1 == r2 and g1 == g2 and b1 == b2
 
 if __name__ == '__main__':
     image = Image.open('')
     no_alpha_image = image.convert('RGB')
     draw = ImageDraw.Draw(no_alpha_image)
-    for region in find_regions(no_alpha_image):
+    for region in find_regions(no_alpha_image, [0,255,0]):
         draw.rectangle(expand_bounding(no_alpha_image, region), outline=(0, 255, 0))
     no_alpha_image.show()