From 9847a7f33f0d3d7e3fe493455bd5b088077782bc Mon Sep 17 00:00:00 2001 From: SoftArmpit Date: Fri, 15 Jun 2018 15:31:00 +0100 Subject: [PATCH] feat: pass whole images and auto-detect shapes We add auto-detection of mask shapes, enabling the ability to put whole images in the input folder. This is done through OpenCV by checking for large contours (>150) matching the mask color. Depends: https://github.com/deeppomf/DeepMindBreak/pull/8 --- Dockerfile-py3 | 5 +- decensor.py | 68 +++++++++++++++++++----- requirements.txt | 1 + shape_detect.py | 131 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 191 insertions(+), 14 deletions(-) create mode 100755 shape_detect.py diff --git a/Dockerfile-py3 b/Dockerfile-py3 index 474bdf3..5720b82 100644 --- a/Dockerfile-py3 +++ b/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 diff --git a/decensor.py b/decensor.py index db70365..403095e 100644 --- a/decensor.py +++ b/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) \ No newline at end of file + decensor(args) diff --git a/requirements.txt b/requirements.txt index cd1b7ad..69950b2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ tqdm scipy pyamg matplotlib +opencv-python diff --git a/shape_detect.py b/shape_detect.py new file mode 100755 index 0000000..2736e35 --- /dev/null +++ b/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()