82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
# HED batch processing script; modified from https://github.com/s9xie/hed/blob/master/examples/hed/HED-tutorial.ipynb
|
|
# Step 1: download the hed repo: https://github.com/s9xie/hed
|
|
# Step 2: download the models and protoxt, and put them under {caffe_root}/examples/hed/
|
|
# Step 3: put this script under {caffe_root}/examples/hed/
|
|
# Step 4: run the following script:
|
|
# python batch_hed.py --images_dir=/data/to/path/photos/ --hed_mat_dir=/data/to/path/hed_mat_files/
|
|
# The code sometimes crashes after computation is done. Error looks like "Check failed: ... driver shutting down". You can just kill the job.
|
|
# For large images, it will produce gpu memory issue. Therefore, you better resize the images before running this script.
|
|
# Step 5: run the MATLAB post-processing script "PostprocessHED.m"
|
|
|
|
|
|
import caffe
|
|
import numpy as np
|
|
from PIL import Image
|
|
import os
|
|
import argparse
|
|
import sys
|
|
import scipy.io as sio
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description='batch proccesing: photos->edges')
|
|
parser.add_argument('--caffe_root', dest='caffe_root', help='caffe root', default='../../', type=str)
|
|
parser.add_argument('--caffemodel', dest='caffemodel', help='caffemodel', default='./hed_pretrained_bsds.caffemodel', type=str)
|
|
parser.add_argument('--prototxt', dest='prototxt', help='caffe prototxt file', default='./deploy.prototxt', type=str)
|
|
parser.add_argument('--images_dir', dest='images_dir', help='directory to store input photos', type=str)
|
|
parser.add_argument('--hed_mat_dir', dest='hed_mat_dir', help='directory to store output hed edges in mat file', type=str)
|
|
parser.add_argument('--border', dest='border', help='padding border', type=int, default=128)
|
|
parser.add_argument('--gpu_id', dest='gpu_id', help='gpu id', type=int, default=1)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
args = parse_args()
|
|
for arg in vars(args):
|
|
print('[%s] =' % arg, getattr(args, arg))
|
|
# Make sure that caffe is on the python path:
|
|
caffe_root = args.caffe_root # this file is expected to be in {caffe_root}/examples/hed/
|
|
sys.path.insert(0, caffe_root + 'python')
|
|
|
|
|
|
if not os.path.exists(args.hed_mat_dir):
|
|
print('create output directory %s' % args.hed_mat_dir)
|
|
os.makedirs(args.hed_mat_dir)
|
|
|
|
imgList = os.listdir(args.images_dir)
|
|
nImgs = len(imgList)
|
|
print('#images = %d' % nImgs)
|
|
|
|
caffe.set_mode_gpu()
|
|
caffe.set_device(args.gpu_id)
|
|
# load net
|
|
net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)
|
|
# pad border
|
|
border = args.border
|
|
|
|
for i in range(nImgs):
|
|
if i % 500 == 0:
|
|
print('processing image %d/%d' % (i, nImgs))
|
|
im = Image.open(os.path.join(args.images_dir, imgList[i]))
|
|
|
|
in_ = np.array(im, dtype=np.float32)
|
|
in_ = np.pad(in_, ((border, border), (border, border), (0, 0)), 'reflect')
|
|
|
|
in_ = in_[:, :, 0:3]
|
|
in_ = in_[:, :, ::-1]
|
|
in_ -= np.array((104.00698793, 116.66876762, 122.67891434))
|
|
in_ = in_.transpose((2, 0, 1))
|
|
# remove the following two lines if testing with cpu
|
|
|
|
# shape for input (data blob is N x C x H x W), set data
|
|
net.blobs['data'].reshape(1, *in_.shape)
|
|
net.blobs['data'].data[...] = in_
|
|
# run net and take argmax for prediction
|
|
net.forward()
|
|
fuse = net.blobs['sigmoid-fuse'].data[0][0, :, :]
|
|
# get rid of the border
|
|
fuse = fuse[(border + 35):(-border + 35), (border + 35):(-border + 35)]
|
|
# save hed file to the disk
|
|
name, ext = os.path.splitext(imgList[i])
|
|
sio.savemat(os.path.join(args.hed_mat_dir, name + '.mat'), {'edge_predict': fuse})
|