Chris
8 lat temu
3 zmienionych plików z 270 dodań i 114 usunięć
-
+158 -0process_pixel_labels.py
-
+0 -114test.py
-
+112 -0test_pretrained.py
+ 158
- 0
process_pixel_labels.py
Wyświetl plik
| @@ -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) | |||
+ 0
- 114
test.py
Wyświetl plik
| @@ -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) | |||
+ 112
- 0
test_pretrained.py
Wyświetl plik
| @@ -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) | |||