diff --git a/test.py b/test.py new file mode 100644 index 0000000..e09957f --- /dev/null +++ b/test.py @@ -0,0 +1,114 @@ +#!/usr/bin/python + +import os +import sys +import argparse +import numpy as np +import caffe +import cv2 +import random + + + +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) + + +def presolve(net, args): + net.blobs["data"].reshape(1, 3 if args.color else 1, args.image_size, args.image_size) + net.blobs["gt"].reshape(1, 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()) + 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.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)) + raw = (255 * predict(net, _input, 'out', args)).astype(np.uint8) + + if args.out_dir: + out_fn = os.path.join(args.out_dir, f.replace('/','_')[:-4] + "_raw.png") + cv2.imwrite(out_fn, raw) + + post, coords = post_process(raw) + for idx2 in [1, 2, 3, 0]: + fd.write('%d,%d,' % (width * coords[idx2][0] / 256., height * coords[idx2][1] / 256.)) + fd.write('\n') + + if args.out_dir: + 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("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="Images to predict") + parser.add_argument("out_file", + help="output file listing quad regions") + parser.add_argument("--out-dir", default='', 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("--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) + +