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