Merge pull request #9 from SoftArmpit/automatic-shape-detection

feat: pass whole images and auto-detect shapes

Dockerfile-py3

@@ -1,6 +1,9 @@
 FROM tensorflow/tensorflow:latest-py3
 
-RUN apt update && apt install -y python3 python3-pip python3-tk python3-numpy && apt clean
+RUN apt-get update \
+  && apt-get install -y python3 python3-pip python3-tk python3-numpy libsm6 libxext6 libxrender-dev \
+  && apt-get clean
+RUN apt-get install -y vim
 
 WORKDIR /app
 COPY . /app

decensor.py

@@ -4,12 +4,14 @@ from PIL import Image
 import tqdm
 import os
 import matplotlib.pyplot as plt
+import stat
 import sys
 sys.path.append('..')
 
 from model import Model
 from poisson_blend import blend
 from config import *
+import shape_detect as sd
 
 #TODO: allow variable batch sizes when decensoring. changing BATCH_SIZE will likely result in crashing
 BATCH_SIZE = 1
@@ -17,7 +19,46 @@ BATCH_SIZE = 1
 mask_color = [args.mask_color_red, args.mask_color_green, args.mask_color_blue]
 poisson_blending_enabled = False
 
+def is_file(file):
+    try:
+        return not stat.S_ISDIR(os.stat(file).st_mode)
+    except:
+        return False
+
+def get_files(dir):
+    all_files = os.listdir(dir)
+    filtered_files = list(filter(lambda file: is_file(os.path.join(dir, file)), all_files))
+    return filtered_files
+
+def find_censor_boxes(image_path):
+    (image, boxes) = sd.process_image_path(image_path, tuple(mask_color))
+
+    i = 0
+    for (box_image, cx, cy) in boxes:
+        pil_box_image = sd.cv_to_pillow(box_image)
+        boxes[i] = (pil_box_image, cx, cy)
+        i += 1
+
+    # boxes = map(lambda box: (sd.cv_to_pillow(box[0]), box[1], box[2]), boxes)
+    return (image, boxes)
+
 def decensor(args):
+    subdir = args.decensor_input_path
+    files = sorted(get_files(subdir))
+
+    for file in files:
+        file_path = os.path.join(subdir, file)
+        if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png":
+            print(file_path)
+            (image, boxes) = find_censor_boxes(file_path)
+            decensored_boxes = decensor_boxes(args, boxes)
+            for (box_pillow_image, cx, cy) in decensored_boxes:
+                box_image = sd.pillow_to_cv(box_pillow_image)
+                image = sd.insert_box((box_image, cx, cy), image)
+
+            sd.write_to_file(image, os.path.join(args.decensor_output_path, file))
+
+def decensor_boxes(args, boxes):
     x = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, args.input_channel_size])
     mask = tf.placeholder(tf.float32, [BATCH_SIZE, args.input_size, args.input_size, 1])
     local_x = tf.placeholder(tf.float32, [BATCH_SIZE, args.local_input_size, args.local_input_size, args.input_channel_size])
@@ -33,21 +74,20 @@ def decensor(args):
     saver = tf.train.Saver()
     saver.restore(sess, './models/latest')
 
-    x_decensor = []
     mask_decensor = []
-    for subdir, dirs, files in sorted(os.walk(args.decensor_input_path)):
-        for file in sorted(files):
-            file_path = os.path.join(subdir, file)
-            if os.path.isfile(file_path) and os.path.splitext(file_path)[1] == ".png":
-                print(file_path)
-                image = Image.open(file_path).convert('RGB')
-                image = np.array(image)
-                image = np.array(image / 127.5 - 1)
-                x_decensor.append(image)
+    x_decensor = []
+
+    for (box_image, cx, cy) in boxes:
+        image = np.array(box_image)
+        image = np.array(image / 127.5 - 1)
+        x_decensor.append(image)
+
     x_decensor = np.array(x_decensor)
     print(x_decensor.shape)
     step_num = int(len(x_decensor) / BATCH_SIZE)
 
+    results = []
+
     cnt = 0
     for i in tqdm.tqdm(range(step_num)):
         x_batch = x_decensor[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
@@ -61,10 +101,12 @@ def decensor(args):
             if (poisson_blending_enabled):
                 img = blend(original, img, mask_batch[0,:,:,0])
             output = Image.fromarray(img.astype('uint8'), 'RGB')
-            dst = args.decensor_output_path + '{}.png'.format("{0:06d}".format(cnt))
-            output.save(dst)
+            results.append((output, boxes[cnt][1], boxes[cnt][2]))
             cnt += 1
 
+    tf.reset_default_graph()
+    return results
+
 def get_mask(x_batch):
     points = []
     mask = []
@@ -82,4 +124,4 @@ def get_mask(x_batch):
 if __name__ == '__main__':
     if not os.path.exists(args.decensor_output_path):
         os.makedirs(args.decensor_output_path)
-    decensor(args)
+    decensor(args)

requirements.txt

@@ -3,3 +3,4 @@ tqdm
 scipy
 pyamg
 matplotlib
+opencv-python

shape_detect.py

@@ -0,0 +1,131 @@
+"""Shape detection.
+
+Detect rectangle shapes in images, cut out 128px
+surrounding shape and then pass it to the decensoring
+program and replace the censored tile with the
+decensored one.
+"""
+
+import numpy as np
+import cv2
+import argparse
+from PIL import Image
+import os
+
+
+isExec = __name__ == '__main__'
+
+
+box_size = 128
+def cv_to_pillow(image, i = 0):
+    """
+    converted = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    conv_array = np.array(converted)
+    pil_image = Image.fromarray(conv_array)
+    print(converted)
+    print(conv_array)
+    print(pil_image)
+    return pil_image
+    """
+
+    # TODO(SoftArmpit): This is inefficient, convert directly instead.
+    file_path = os.path.join('/tmp/', str(i) + '.png')
+    write_to_file(image, file_path)
+    pil_box_image = Image.open(file_path).convert('RGB')
+    return pil_box_image
+
+
+def pillow_to_cv(image):
+    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+
+def insert_box(box, image):
+    (box_image, x, y) = box
+    image[y:y+box_image.shape[0], x:x+box_image.shape[1]] = box_image
+    return image
+
+
+def detect_shape(c):
+    perim = cv2.arcLength(c, True)
+    vertices = cv2.approxPolyDP(c, 0.04 * perim, True)
+
+    print('Vertices: ' + str(len(vertices)))
+    return len(vertices) == 4
+
+
+def process_contour(image, c):
+    M = cv2.moments(c)
+    print(M)
+
+    if M['m00'] == 0:
+        return None
+
+    cx = int(M['m10'] / M['m00'] - box_size / 2)
+    cy = int(M['m01'] / M['m00'] - box_size / 2)
+    # NOTE(SoftArmpit): Limit box to image boundaries
+    cx = min(max(cx, 0), image.shape[1] - box_size)
+    cy = min(max(cy, 0), image.shape[0] - box_size)
+    box = image[cy:cy+box_size, cx:cx+box_size]
+
+    print(str(cx) + ", " + str(cy))
+    area = cv2.contourArea(c)
+
+    if area < 148:
+        print('Area too small: ' + str(area) + "at " + str(cx) + 'x' + str(cy))
+        return None
+
+    return (box, cx, cy)
+
+
+def process_image(image, mask_color):
+    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+    green_mask = cv2.inRange(image, mask_color, mask_color)
+    if isExec:
+        cv2.imshow("Mask", green_mask)
+    (_, cs, _) = cv2.findContours(green_mask,
+                                  cv2.RETR_EXTERNAL,
+                                  cv2.CHAIN_APPROX_SIMPLE)
+
+    boxes = []
+    for c in cs:
+        isRect = detect_shape(c)
+
+        if isRect or True:
+            print("Rectangle detected")
+            pc = process_contour(image, c)
+            if pc is not None:
+                boxes.append(pc)
+
+    return boxes
+
+
+def process_image_path(image_path, mask_color):
+    image = cv2.imread(image_path)
+    return (image, process_image(image, mask_color))
+
+
+def write_to_file(image, path):
+    cv2.imwrite(path, image)
+
+
+def main():
+    """Entry function."""
+    ap = argparse.ArgumentParser()
+    ap.add_argument('-f', '--file', required=True, help='Path to image file')
+    ap.add_argument('-g', '--green', required=False, default=255)
+    ap.add_argument('-r', '--red', required=False, default=0)
+    ap.add_argument('-b', '--blue', required=False, default=0)
+    args = ap.parse_args()
+
+    (image, boxes) = process_image_path(args.file, (args.red, args.green, args.blue))
+
+    print(len(boxes))
+    for (box_image, cx, cy) in boxes:
+        cv2.imshow("Box " + str(cx) + 'x' + str(cy), box_image)
+
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()
+
+
+if __name__ == '__main__':
+    main()