From 83d587c56ef6418d1377723304b1f843b86b1251 Mon Sep 17 00:00:00 2001 From: Chris Date: Tue, 12 Sep 2017 09:24:59 -0600 Subject: [PATCH] Post processing --- process_pixel_labels.py | 158 ++++++++++++++++++++++++++++++++++++++++ test.py | 114 ----------------------------- test_pretrained.py | 112 ++++++++++++++++++++++++++++ 3 files changed, 270 insertions(+), 114 deletions(-) create mode 100644 process_pixel_labels.py delete mode 100644 test.py create mode 100644 test_pretrained.py diff --git a/process_pixel_labels.py b/process_pixel_labels.py new file mode 100644 index 0000000..4615f17 --- /dev/null +++ b/process_pixel_labels.py @@ -0,0 +1,158 @@ +import cv2 +import os +import numpy as np +import sys + +def draw_poly(img, bounding_poly): + pts = np.array(bounding_poly, np.int32) + + #http://stackoverflow.com/a/15343106/3479446 + mask = np.zeros(img.shape[:2], dtype=np.uint8) + roi_corners = np.array([pts], dtype=np.int32) + + ignore_mask_color = (255,) + cv2.fillPoly(mask, roi_corners, ignore_mask_color, lineType=cv2.LINE_8) + return mask + +def post_process(img): + # img = open_close(img) + img = get_largest_cc(img) + img = fill_holes(img) + + # img = min_area_rectangle(img) + img, coords = improve_min_area_rectangle(img) + + return img, coords + +def open_close(img): + kernel = np.ones((3,3),np.uint8) + erosion = cv2.erode(img,kernel,iterations = 15) + dilation = cv2.dilate(erosion,kernel,iterations = 15) + + return dilation + + +def get_largest_cc(img): + img = img.copy() + ret, th = cv2.threshold(img,127,255,cv2.THRESH_BINARY) + connectivity = 4 + output= cv2.connectedComponentsWithStats(th, connectivity, cv2.CV_32S) + cnts = output[2][1:,4] + largest = cnts.argmax() + 1 + img[output[1] != largest] = 0 + + return img + +def get_iou(gt_img, pred_img): + inter = gt_img & pred_img + union = gt_img | pred_img + + iou = np.count_nonzero(inter) / float(np.count_nonzero(union)) + + return iou + +def draw_box(img, box): + box = np.int0(box) + draw = np.zeros_like(img) + cv2.drawContours(draw,[box],0,(255),-1) + return draw + +def compute_iou(img, box): + # box = np.int0(box) + # draw = np.zeros_like(img) + # cv2.drawContours(draw,[box],0,(255),-1) + draw = draw_box(img, box) + v = get_iou(img, draw) + return v + +def step_box(img, box, step_size=1): + best_val = -1 + best_box = None + for index, x in np.ndenumerate(box): + for d in [-step_size, step_size]: + alt_box = box.copy() + alt_box[index] = x + d + + v = compute_iou(img, alt_box) + if best_val < v: + best_val = v + best_box = alt_box + return best_val, best_box + + + +def improve_min_area_rectangle(img): + img = img.copy() + _, contours,_ = cv2.findContours(img, 1, 2) + cnt = contours[0] + rect = cv2.minAreaRect(cnt) + box = cv2.boxPoints(rect) + + best_val = compute_iou(img, box) + best_box = box + + while True: + new_val, new_box = step_box(img, best_box, step_size=1) + # print new_val + if new_val <= best_val: + break + best_val = new_val + best_box = new_box + + return draw_box(img, best_box), best_box + + +def min_area_rectangle(img): + img = img.copy() + _, contours,_ = cv2.findContours(img, 1, 2) + cnt = contours[0] + + rect = cv2.minAreaRect(cnt) + box = cv2.boxPoints(rect) + box = np.int0(box) + draw = np.zeros_like(img) + cv2.drawContours(draw,[box],0,(255),-1) + + return draw + + +def fill_holes(img): + im_th = img.copy() + + + # Copy the thresholded image. + im_floodfill = im_th.copy() + + # Mask used to flood filling. + # Notice the size needs to be 2 pixels than the image. + h, w = im_th.shape[:2] + mask = np.zeros((h+2, w+2), np.uint8) + + # Floodfill from point (0, 0) + if img[0,0] != 0: + print "WARNING: Filling something you shouldn't" + cv2.floodFill(im_floodfill, mask, (0,0), 255); + + # Invert floodfilled image + im_floodfill_inv = cv2.bitwise_not(im_floodfill) + + # Combine the two images to get the foreground. + im_out = im_th | im_floodfill_inv + + return im_out + +if __name__ == "__main__": + pred_folder = sys.argv[1] + out_folder = sys.argv[2] + + pred_imgs = {} + for root, folders, files in os.walk(pred_folder): + for f in files: + if f.endswith(".png"): + pred_imgs[f] = os.path.join(root, f) + + for k in pred_imgs: + pred_img = cv2.imread(pred_imgs[k], 0) + post_img = post_process(pred_img) + + cv2.imwrite(os.path.join(out_folder, k), post_img) diff --git a/test.py b/test.py deleted file mode 100644 index 49e4bcc..0000000 --- a/test.py +++ /dev/null @@ -1,114 +0,0 @@ -#!/usr/bin/python - -import os -import sys -import collections -import argparse -import numpy as np -import matplotlib -matplotlib.use("AGG") -import matplotlib.pyplot as plt -import caffe -import cv2 -import random -import scipy.ndimage as nd - - - -def safe_mkdir(_dir): - try: - os.makedirs(_dir) - except: - pass - - -def predict(network, im, output_blob, args): - network.blobs["data"].data[0,:,:,:] = im - network.forward() - - #response = network.blobs[output_blob].data[0,:].copy() - #return np.argmax(response, axis=0) - - response = network.blobs[output_blob].data[0,0,:,:].copy() - response[response >= 0.5] = 1 - response[response <= 0.5] = 0 - return response - - -def presolve(net, args): - net.blobs["data"].reshape(args.batch_size, 3 if args.color else 1, args.image_size, args.image_size) - net.blobs["gt"].reshape(args.batch_size, 1, args.image_size, args.image_size) - - -def main(args): - net = caffe.Net(args.net_file, args.weight_file, caffe.TEST) - presolve(net, args) - - file_list = map(lambda s: s.strip(), open(args.test_manifest, 'r').readlines()) - for idx, line in enumerate(file_list): - if idx % args.print_count == 0: - print "Processed %d/%d Images" % (idx, len(file_list)) - tokens = line.split(',') - f = tokens[0] - resolved = os.path.join(args.dataset_dir, f) - im = cv2.imread(resolved, 1 if args.color else 0) - - _input = args.scale * (cv2.resize(im, (args.image_size, args.image_size)) - args.mean) - if _input.ndim > 2: - _input = np.transpose(_input, (2, 0, 1)) - output = predict(net, _input, 'out', args) - - out_fn = os.path.join(args.out_dir, f.replace('/','_')[:-4] + ".png") - cv2.imwrite(out_fn, (255 * output).astype(np.uint8)) - - -def get_args(): - parser = argparse.ArgumentParser(description="Outputs binary predictions") - - parser.add_argument("net_file", - help="The deploy.prototxt") - parser.add_argument("weight_file", - help="The .caffemodel") - parser.add_argument("dataset_dir", - help="The dataset to be evaluated") - parser.add_argument("test_manifest", - help="txt file listing images to train on") - parser.add_argument("--out-dir", default='out', type=str, - help="Dump images") - - parser.add_argument("--gpu", type=int, default=0, - help="GPU to use for running the network") - parser.add_argument("-c", "--color", default=False, action='store_true', - help="Training batch size") - - parser.add_argument("-m", "--mean", type=float, default=127., - help="Mean value for data preprocessing") - parser.add_argument("-s", "--scale", type=float, default=0.0039, - help="Optional pixel scale factor") - parser.add_argument("-b", "--batch-size", default=2, type=int, - help="Training batch size") - parser.add_argument("--image-size", default=256, type=int, - help="Size of images for input to prediction") - - parser.add_argument("--print-count", default=10, type=int, - help="Print interval") - - args = parser.parse_args() - print args - - return args - - -if __name__ == "__main__": - args = get_args() - safe_mkdir(args.out_dir) - - if args.gpu >= 0: - caffe.set_device(args.gpu) - caffe.set_mode_gpu() - else: - caffe.set_mode_cpu() - - main(args) - - diff --git a/test_pretrained.py b/test_pretrained.py new file mode 100644 index 0000000..eabc408 --- /dev/null +++ b/test_pretrained.py @@ -0,0 +1,112 @@ +#!/usr/bin/python + +import os +import sys +import argparse +import numpy as np +import caffe +import cv2 +from process_pixel_labels import post_process + +NET_FILE = './models/cbad_train_val.prototxt' +WEIGHT_FILE = './models/cbad_weights.prototxt' + + +def safe_mkdir(_dir): + try: + os.makedirs(_dir) + except: + pass + + +def predict(network, im, output_blob, args): + network.blobs["data"].data[0,:,:,:] = im + network.forward() + + if args.model == 'ohio': + # sigmoid + response = network.blobs[output_blob].data[0,0,:,:].copy() + response[response >= 0.5] = 1 + response[response <= 0.5] = 0 + return response + else: + # softmax + response = network.blobs[output_blob].data[0,:].copy() + return np.argmax(response, axis=0) + + + +def presolve(net, args): + net.blobs["data"].reshape(args.batch_size, 3, args.image_size, args.image_size) + net.blobs["gt"].reshape(args.batch_size, 1, args.image_size, args.image_size) + + +def main(args): + net = caffe.Net(NET_FILE, WEIGHT_FILE, caffe.TEST) + presolve(net, args) + + file_list = map(lambda s: s.strip(), open(args.manifest, 'r').readlines()) + fd = open(args.out_file, 'w') + for idx, line in enumerate(file_list): + if idx % args.print_count == 0: + print "Processed %d/%d Images" % (idx, len(file_list)) + tokens = line.split(',') + f = tokens[0] + resolved = os.path.join(args.image_dir, f) + im = cv2.imread(resolved, 1) + + _input = 0.0039 * (cv2.resize(im, (256, 256)) - 127.) + _input = np.transpose(_input, (2, 0, 1)) + raw = (255 * predict(net, _input, 'out', args)).astype(np.uint8) + + out_fn = os.path.join(args.out_dir, f.replace('/','_')[:-4] + "_raw.png") + cv2.imwrite(out_fn, raw) + + post, coords = post_process(raw) + + out_fn = os.path.join(args.out_dir, f.replace('/','_')[:-4] + "_post.png") + cv2.imwrite(out_fn, post) + + +def get_args(): + parser = argparse.ArgumentParser(description="Outputs binary predictions") + + parser.add_argument("image_dir", + help="The directory where images are stored") + parser.add_argument("manifest", + help="txt file listing images relative to image_dir") + parser.add_argument("model", + help="[cbad|ohio]") + parser.add_argument("out_file", type=str, + help="Output file") + + parser.add_argument("--out-dir", type=str, default=out, + help="") + parser.add_argument("--gpu", type=int, default=0, + help="GPU to use for running the network") + parser.add_argument("--print-count", default=10, type=int, + help="Print interval") + + args = parser.parse_args() + print args + + return args + + +if __name__ == "__main__": + args = get_args() + safe_mkdir(args.out_dir) + + if args.model == 'ohio': + NET_FILE = './models/ohio_train_val.prototxt' + WEIGHT_FILE = './models/ohio_weights.caffemodel' + + if args.gpu >= 0: + caffe.set_device(args.gpu) + caffe.set_mode_gpu() + else: + caffe.set_mode_cpu() + + main(args) + +