Post processing

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)

test.py

@@ -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)
-
-

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)
+
+