127 lines
3.3 KiB
Python
127 lines
3.3 KiB
Python
import cv2
|
|
import base64
|
|
import os
|
|
import sys
|
|
import copy
|
|
|
|
sys.path.insert(0, ".")
|
|
|
|
from tools.infer.utility import parse_args
|
|
from tools.infer.predict_det import TextDetector
|
|
|
|
|
|
def cv2_to_base64(image):
|
|
return base64.b64encode(image).decode('utf8')
|
|
|
|
def merge_configs():
|
|
# deafult cfg
|
|
backup_argv = copy.deepcopy(sys.argv)
|
|
sys.argv = sys.argv[:1]
|
|
cfg = parse_args()
|
|
|
|
update_cfg_map = vars(read_params())
|
|
|
|
for key in update_cfg_map:
|
|
cfg.__setattr__(key, update_cfg_map[key])
|
|
|
|
sys.argv = copy.deepcopy(backup_argv)
|
|
return cfg
|
|
|
|
class Config(object):
|
|
pass
|
|
|
|
|
|
def read_params():
|
|
cfg = Config()
|
|
|
|
# params for text detector
|
|
cfg.det_algorithm = "DB"
|
|
# cfg.det_model_dir = "./inference/idCard2/ch_db_mv3_sfz/"
|
|
cfg.det_model_dir = "./output_inference/ch_db_mv3_sfz/"
|
|
# cfg.det_model_dir = "./inference/idCard2/det_r50_vd_db_sfz/"
|
|
cfg.det_resize_long = 960
|
|
|
|
# DB parmas
|
|
cfg.det_db_thresh = 0.3
|
|
cfg.det_db_box_thresh = 0.5
|
|
cfg.det_db_unclip_ratio = 2.0
|
|
cfg.use_dilation = False
|
|
cfg.det_db_score_mode = "fast"
|
|
|
|
# #EAST parmas
|
|
# cfg.det_east_score_thresh = 0.8
|
|
# cfg.det_east_cover_thresh = 0.1
|
|
# cfg.det_east_nms_thresh = 0.2
|
|
|
|
cfg.use_pdserving = False
|
|
cfg.use_tensorrt = False
|
|
|
|
return cfg
|
|
|
|
|
|
# 身份证识别
|
|
class IdCard(object):
|
|
# 初始化
|
|
def __init__(self, text_system, args):
|
|
self.text_system = text_system
|
|
self.score = 0.5
|
|
# 实例分割
|
|
# if args["idCard"]["accuracy"] == "accurate":
|
|
# self.directory = "./inference/idCard/cascade_mask_rcnn_mobilenetv3_fpn_1x/"
|
|
# else:
|
|
# self.directory = "./inference/idCard/mask_rcnn_mobilenetv3_fpn_1x/"
|
|
|
|
cfg = merge_configs()
|
|
|
|
cfg.use_gpu = args["use_gpu"]
|
|
if args["use_gpu"]:
|
|
try:
|
|
_places = os.environ["CUDA_VISIBLE_DEVICES"]
|
|
int(_places[0])
|
|
print("use gpu: ", args["use_gpu"])
|
|
print("CUDA_VISIBLE_DEVICES: ", _places)
|
|
cfg.gpu_mem = 8000
|
|
except:
|
|
raise RuntimeError(
|
|
"Environment Variable CUDA_VISIBLE_DEVICES is not set correctly. If you wanna use gpu, please set CUDA_VISIBLE_DEVICES via export CUDA_VISIBLE_DEVICES=cuda_device_id."
|
|
)
|
|
cfg.ir_optim = True
|
|
cfg.enable_mkldnn = args["enable_mkldnn"]
|
|
|
|
self.model = TextDetector(cfg)
|
|
|
|
def idCard(self, img):
|
|
dt_boxes, dt_labels, elapse = self.model(img, cls=True)
|
|
|
|
print(dt_boxes)
|
|
print(dt_labels)
|
|
new_results = {}
|
|
|
|
categorys = ["name", "sex", "nation", "birthday", "address", "number"]
|
|
|
|
if str(new_results) == '{}':
|
|
return ""
|
|
|
|
return new_results
|
|
|
|
def __call__(self, img=None):
|
|
return self.idCard(img)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# from deploy.hubserving.ocr_system.module import OCRSystem
|
|
# text_system = OCRSystem({
|
|
# "use_gpu": True,
|
|
# "enable_mkldnn": False
|
|
# }).text_sys
|
|
|
|
module = IdCard(None, {
|
|
# "use_gpu": True,
|
|
"use_gpu": False,
|
|
"enable_mkldnn": False
|
|
})
|
|
# image = cv2.imread('E:\\dataset\\shenfenzheng\\images\\JPEGImages\\001.png')
|
|
img_path = './id_cards/card1.png'
|
|
image = cv2.imread(img_path)
|
|
res = module(img=image)
|