def __init__(self, base_path):
self.base_path = base_path
self.nms = NMSWrapper(self.DEFAULT_NMS_TYPE)
self.net = FasterRCNNSlim()
cfg = tf.ConfigProto()
self.sess = tf.Session(config=cfg)
saver = tf.train.Saver()
saver.restore(self.sess, self.DEFAULT_MODEL)
def get_session():
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
sess = tf.Session(config=cfg)
net = FasterRCNNSlim()
saver = tf.train.Saver()
saver.restore(sess, "model/res101_faster_rcnn_iter_60000.ckpt")
return (sess, net)
def load(self, base):
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
self.sess = tf.Session(config=cfg)
self.net = FasterRCNNSlim()
saver = tf.train.Saver()
path = os.path.join(base, "animeFace/model/res101_faster_rcnn_iter_60000.ckpt")
try:
saver.restore(self.sess, path)
except: print("animeFace:", path); return False
return True
def __init__(self,
model_path=None,
nms_type='CPU_NMS',
nms_threshold=0.3,
threshold=0.8):
"""
Parameters:
nms_type: Type of NMS. Options ("PY_NMS" | "CPU_NMS" | "GPU_NMS")
nms_threshold: Threshold for Non Max Suppression. (float)
threshold: Threshold for class regression. (float)
model_path: Path to the model. If none is provided, it uses the one inside this directory.
quiet: Disable tensorflow display messages. (boolean)
"""
# nms_type check
if nms_type == 'PY_NMS':
nms_type = NMSType.PY_NMS
elif nms_type == 'CPU_NMS':
nms_type = NMSType.CPU_NMS
elif nms_type == 'GPU_NMS':
nms_type = NMSType.GPU_NMS
else:
raise ValueError('Incorrect NMS Type, not supported yet')
self.nms = NMSWrapper(nms_type)
self.nms_threshold = nms_threshold
self.cls_threshold = threshold
cfg = tf.compat.v1.ConfigProto()
cfg.gpu_options.allow_growth = True
self.sess = tf.compat.v1.Session(config=cfg)
self.net = FasterRCNNSlim()
saver = tf.compat.v1.train.Saver()
if model_path is None:
model_path = str(
Path(__file__).resolve().parent /
'model/res101_faster_rcnn_iter_60000.ckpt')
saver.restore(self.sess, model_path)
def main():
parser = argparse.ArgumentParser(description='Anime face detector demo')
parser.add_argument('-i',
help='The input path of an image or directory',
required=True,
dest='input',
type=str)
parser.add_argument('-o',
help='The output json path of the detection result',
dest='output')
parser.add_argument(
'-nms',
help='Change the threshold for non maximum suppression',
dest='nms_thresh',
default=0.3,
type=float)
parser.add_argument('-conf',
help='Change the threshold for class regression',
dest='conf_thresh',
default=0.8,
type=float)
parser.add_argument('-model',
help='Specify a new path for model',
dest='model',
type=str,
default='model/res101_faster_rcnn_iter_60000.ckpt')
parser.add_argument('-nms-type',
help='Type of nms',
choices=['PY_NMS', 'CPU_NMS', 'GPU_NMS'],
dest='nms_type',
default='CPU_NMS')
args = parser.parse_args()
assert os.path.exists(args.input), 'The input path does not exists'
if os.path.isdir(args.input):
files = load_file_from_dir(args.input)
else:
files = [args.input]
file_len = len(files)
if args.nms_type == 'PY_NMS':
nms_type = NMSType.PY_NMS
elif args.nms_type == 'CPU_NMS':
nms_type = NMSType.CPU_NMS
elif args.nms_type == 'GPU_NMS':
nms_type = NMSType.GPU_NMS
else:
raise ValueError('Incorrect NMS Type, not supported yet')
nms = NMSWrapper(nms_type)
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
sess = tf.Session(config=cfg)
net = FasterRCNNSlim()
saver = tf.train.Saver()
saver.restore(sess, args.model)
result = {}
time_start = time.time()
for idx, file in enumerate(files):
elapsed = time.time() - time_start
eta = (file_len - idx) * elapsed / idx if idx > 0 else 0
print('[%d/%d] Elapsed: %s, ETA: %s >> %s' %
(idx + 1, file_len, fmt_time(elapsed), fmt_time(eta), file))
img = cv2.imread(file)
scores, boxes = detect(sess, net, img)
boxes = boxes[:, 4:8]
scores = scores[:, 1]
keep = nms(
np.hstack([boxes, scores[:, np.newaxis]]).astype(np.float32),
args.nms_thresh)
boxes = boxes[keep, :]
scores = scores[keep]
inds = np.where(scores >= args.conf_thresh)[0]
scores = scores[inds]
boxes = boxes[inds, :]
result[file] = []
for i in range(scores.shape[0]):
x1, y1, x2, y2 = boxes[i, :].tolist()
new_result = {'score': float(scores[i]), 'bbox': [x1, y1, x2, y2]}
result[file].append(new_result)
if args.output is None:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)),
(0, 0, 255), 2)
if args.output:
if ((idx + 1) % 1000) == 0:
# saving the temporary result
with open(args.output, 'w') as f:
json.dump(result, f)
else:
cv2.imshow(file, img)
if args.output:
with open(args.output, 'w') as f:
json.dump(result, f)
else:
cv2.waitKey()
def main():
parser = argparse.ArgumentParser(description='Anime face detector demo')
parser.add_argument('-i',
help='The input path of an image or directory',
required=True,
dest='input',
type=str)
parser.add_argument('-o',
help='The output json path of the detection result',
dest='output')
parser.add_argument(
'-nms',
help='Change the threshold for non maximum suppression',
dest='nms_thresh',
default=0.3,
type=float)
parser.add_argument('-conf',
help='Change the threshold for class regression',
dest='conf_thresh',
default=0.8,
type=float)
parser.add_argument('-model',
help='Specify a new path for model',
dest='model',
type=str,
default='model/res101_faster_rcnn_iter_60000.ckpt')
parser.add_argument('-nms-type',
help='Type of nms',
choices=['PY_NMS', 'CPU_NMS', 'GPU_NMS'],
dest='nms_type',
default='CPU_NMS')
parser.add_argument('-crop-location',
help='The output folder to place the cropped images',
dest='crop_output_image_location')
parser.add_argument(
'-start-output',
help='Start the numbering of the cropped images filename',
dest='start_output_number',
default=0,
type=int)
parser.add_argument('-crop-width',
help='The width of images to crop',
dest='crop_width',
type=int)
parser.add_argument('-crop-height',
help='The height of images to crop',
dest='crop_height',
type=int)
parser.add_argument('-scale-factor',
help='Scale factor for the bounding boxes',
dest='scale_factor',
type=float)
parser.add_argument('-square-crop', dest='square', action='store_true')
parser.add_argument(
'-min-len',
help='The minimum length of the smallest side of the bounding box',
dest='min_len',
type=int)
parser.set_defaults(square=False)
args = parser.parse_args()
assert os.path.exists(args.input), 'The input path does not exists'
if os.path.isdir(args.input):
files = load_file_from_dir(args.input)
else:
files = [args.input]
file_len = len(files)
if args.nms_type == 'PY_NMS':
nms_type = NMSType.PY_NMS
elif args.nms_type == 'CPU_NMS':
nms_type = NMSType.CPU_NMS
elif args.nms_type == 'GPU_NMS':
nms_type = NMSType.GPU_NMS
else:
raise ValueError('Incorrect NMS Type, not supported yet')
nms = NMSWrapper(nms_type)
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
sess = tf.Session(config=cfg)
net = FasterRCNNSlim()
saver = tf.train.Saver()
saver.restore(sess, args.model)
result = {}
time_start = time.time()
for idx, file in enumerate(files):
save_name = file.split('/')[-1].split('.')[0]
elapsed = time.time() - time_start
eta = (file_len - idx) * elapsed / idx if idx > 0 else 0
print('[%d/%d] Elapsed: %s, ETA: %s >> %s' %
(idx + 1, file_len, fmt_time(elapsed), fmt_time(eta), file))
img = cv2.imread(file)
scores, boxes = detect(sess, net, img)
boxes = boxes[:, 4:8]
scores = scores[:, 1]
keep = nms(
np.hstack([boxes, scores[:, np.newaxis]]).astype(np.float32),
args.nms_thresh)
boxes = boxes[keep, :]
scores = scores[keep]
inds = np.where(scores >= args.conf_thresh)[0]
scores = scores[inds]
boxes = boxes[inds, :]
result[file] = []
cnt = 1
for i in range(scores.shape[0]):
x1, y1, x2, y2 = boxes[i, :].tolist()
new_result = {'score': float(scores[i]), 'bbox': [x1, y1, x2, y2]}
result[file].append(new_result)
if args.output is None and args.crop_output_image_location is None:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)),
(0, 0, 255), 2)
if args.crop_output_image_location:
width = x2 - x1
height = y2 - y1
center_w = (x2 + x1) / 2
center_h = (y2 + y1) / 2
if args.scale_factor:
width *= args.scale_factor
height *= args.scale_factor
if args.square:
box_len = max(width, height)
x2 = int(center_w + box_len / 2)
x1 = int(center_w - box_len / 2)
y2 = int(center_h + box_len / 2)
y1 = int(center_h - box_len / 2)
else:
x2 = int(center_w + width / 2)
x1 = int(center_w - width / 2)
y2 = int(center_h + height / 2)
y1 = int(center_h - height / 2)
x2 = min(x2, img.width - 1)
y2 = min(y2, img.height - 1)
x1 = max(x1, 0)
y1 = max(y1, 0)
if args.min_len:
if min(width, height) < args.min_len: continue
cropped_image = img[int(y1):int(y2), int(x1):int(x2)]
if args.crop_width and args.crop_height:
cropped_image = cv2.resize(
cropped_image, (args.crop_width, args.crop_height),
interpolation=cv2.INTER_AREA)
if cropped_image.any():
cv2.imwrite(
args.crop_output_image_location + save_name + '_' +
str(cnt) + ".jpg", cropped_image)
args.start_output_number += 1
cnt += 1
else:
continue
if args.output:
if ((idx + 1) % 1000) == 0:
# saving the temporary result
with open(args.output, 'w') as f:
json.dump(result, f)
elif args.crop_output_image_location is None:
cv2.imshow(file, img)
if args.output:
with open(args.output, 'w') as f:
json.dump(result, f)
else:
cv2.waitKey()
def main():
parser = argparse.ArgumentParser(description='Anime face detector demo')
parser.add_argument('-i', help='The input path of an image or directory', required=True, dest='input', type=str)
parser.add_argument('-o', help='The output json path of the detection result', dest='output')
parser.add_argument('-c', help='The output path of faces images that being detected', dest='cut', type=str)
parser.add_argument('-nms', help='Change the threshold for non maximum suppression',
dest='nms_thresh', default=0.3, type=float)
parser.add_argument('-conf', help='Change the threshold for class regression', dest='conf_thresh',
default=0.8, type=float)
parser.add_argument('-model', help='Specify a new path for model', dest='model', type=str,
default='model/res101_faster_rcnn_iter_60000.ckpt')
parser.add_argument('-nms-type', help='Type of nms', choices=['PY_NMS', 'CPU_NMS', 'GPU_NMS'], dest='nms_type',
default='CPU_NMS')
args = parser.parse_args()
assert os.path.exists(args.input), 'The input path does not exists'
if os.path.isdir(args.input):
files = load_file_from_dir(args.input)
else:
files = [args.input]
file_len = len(files)
if args.nms_type == 'PY_NMS':
nms_type = NMSType.PY_NMS
elif args.nms_type == 'CPU_NMS':
nms_type = NMSType.CPU_NMS
elif args.nms_type == 'GPU_NMS':
nms_type = NMSType.GPU_NMS
else:
raise ValueError('Incorrect NMS Type, not supported yet')
nms = NMSWrapper(nms_type)
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
sess = tf.Session(config=cfg)
net = FasterRCNNSlim()
saver = tf.train.Saver()
saver.restore(sess, args.model)
result = {}
time_start = time.time()
recognized_num=1
last_file=''
# Config
configpath = os.path.join(os.path.dirname(sys.argv[0]), 'lastrun.ini')
config = ConfigParser()
if not os.path.exists(configpath):
config['LastRun'] = {
'recognized_num' : recognized_num,
'last_file':''
}
with open(configpath, 'w', encoding='utf-8') as file:
config.write(file)
# check whether path exists.
if not os.path.exists(configpath):
config['LastRun'] = {
'recognized_num' : recognized_num,
'last_file' : file
}
with open(configpath, 'w', encoding='utf-8') as f:
config.write(f)
else:
config.readfp(codecs.open(configpath, 'r', 'utf8'))
recognized_num = config.getint('LastRun', 'recognized_num')
last_file = config.get('LastRun', 'last_file')
if last_file in files:
index = files.index(last_file)
del files[0:index]
file_len = len(files)
# make sure that cut save path exists.
if args.cut is not None:
if not os.path.isdir(os.path.dirname(args.cut)):
os.makedirs(os.path.dirname(args.cut))
for idx, file in enumerate(files):
elapsed = time.time() - time_start
eta = (file_len - idx) * elapsed / idx if idx > 0 else 0
print('[%d/%d] Elapsed: %s, ETA: %s >> %s' % (idx+1, file_len, fmt_time(elapsed), fmt_time(eta), file))
img = cv2.imread(file)
# check whether file is null.
if img is None:
print('[Warning] Can\'t read %s therefore skipped. NOTICE: THIS FILE MAY BE A GIF IMAGE.' % file)
continue
scores, boxes = detect(sess, net, img)
boxes = boxes[:, 4:8]
scores = scores[:, 1]
keep = nms(np.hstack([boxes, scores[:, np.newaxis]]).astype(np.float32), args.nms_thresh)
boxes = boxes[keep, :]
scores = scores[keep]
inds = np.where(scores >= args.conf_thresh)[0]
scores = scores[inds]
boxes = boxes[inds, :]
result[file] = []
for i in range(scores.shape[0]):
x1, y1, x2, y2 = boxes[i, :].tolist()
new_result = {'score': float(scores[i]),
'bbox': [x1, y1, x2, y2]}
result[file].append(new_result)
if args.output is None and args.cut is None:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
elif args.cut is not None:
detected = img[int(y1):int(y2), int(x1):int(x2)]
output = os.path.join(args.cut, str(recognized_num) + '.png')
cv2.imwrite(output, detected)
recognized_num += 1
if args.output:
# save temporary result every 10 files.
if ((idx+1) % 10) == 0:
# saving the temporary result
with open(args.output, 'w') as f:
json.dump(result, f)
elif args.cut is None:
cv2.imshow(file, img)
# save last round status.
config['LastRun'] = {
'recognized_num' : recognized_num,
'last_file' : file
}
with open(configpath, 'w', encoding='utf-8') as f:
config.write(f)
if args.output:
with open(args.output, 'w') as f:
json.dump(result, f)
elif args.cut is None:
cv2.waitKey()
def main():
parser = gen_parser()
args = parser.parse_args()
assert os.path.exists(args.input), 'The input path does not exists'
if os.path.isdir(args.input):
files = load_file_from_dir(args.input)
else:
files = [args.input]
file_len = len(files)
if args.nms_type == 'PY_NMS':
nms_type = NMSType.PY_NMS
elif args.nms_type == 'CPU_NMS':
nms_type = NMSType.CPU_NMS
elif args.nms_type == 'GPU_NMS':
nms_type = NMSType.GPU_NMS
else:
raise ValueError('Incorrect NMS Type, not supported yet')
nms = NMSWrapper(nms_type)
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
sess = tf.Session(config=cfg)
net = FasterRCNNSlim()
saver = tf.train.Saver()
saver.restore(sess, args.model)
result = {}
time_start = time.time()
for idx, file in enumerate(files):
elapsed = time.time() - time_start
eta = (file_len - idx) * elapsed / idx if idx > 0 else 0
print('[%d/%d] Elapsed: %s, ETA: %s >> %s' % (idx+1, file_len, fmt_time(elapsed), fmt_time(eta), file))
img = cv2.imread(file)
scores, boxes = detect(sess, net, img)
boxes = boxes[:, 4:8]
scores = scores[:, 1]
keep = nms(np.hstack([boxes, scores[:, np.newaxis]]).astype(np.float32), args.nms_thresh)
boxes = boxes[keep, :]
scores = scores[keep]
inds = np.where(scores >= args.conf_thresh)[0]
scores = scores[inds]
boxes = boxes[inds, :]
result[file] = []
for i in range(scores.shape[0]):
x1, y1, x2, y2 = boxes[i, :].tolist()
new_result = {'score': float(scores[i]),
'bbox': [x1, y1, x2, y2]}
result[file].append(new_result)
if args.output is None:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
if args.output:
if ((idx+1) % 1000) == 0:
# saving the temporary result
with open(args.output, 'w') as f:
json.dump(result, f)
elif args.stdout:
print(json.dumps(result))
else:
cv2.imshow(file, img)
if args.output:
with open(args.output, 'w') as f:
json.dump(result, f)
elif args.stdout:
pass
else:
cv2.waitKey()