145 lines
4.6 KiB
Python
145 lines
4.6 KiB
Python
import toml
|
|
from loguru import logger
|
|
import zmq
|
|
import cv2
|
|
import numpy as np
|
|
import requests
|
|
import base64
|
|
import datetime
|
|
|
|
|
|
|
|
def get_access_token():
|
|
client_id = "MDCGplPqK0kteOgbXwt5cyn0"
|
|
client_secret = "yIHJQUUiMkkw53nlQqHpiLvRFsLGcqgn"
|
|
url = "https://aip.baidubce.com/oauth/2.0/token"
|
|
params = {
|
|
'grant_type': 'client_credentials',
|
|
'client_id': client_id,
|
|
'client_secret': client_secret
|
|
}
|
|
headers = {
|
|
'Content-Type': 'application/json',
|
|
'Accept': 'application/json'
|
|
}
|
|
response = requests.post(url, params=params, headers=headers)
|
|
response_json = response.json()
|
|
if 'access_token' in response_json:
|
|
return response_json['access_token']
|
|
else:
|
|
print("Failed to get access_token:", response_json.get('error_description'))
|
|
return None
|
|
|
|
def ocr_api_request(image_base64):
|
|
# url = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate_basic" # 高精度
|
|
url = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate" # 高精度带位置
|
|
# url = "https://aip.baidubce.com/rest/2.0/ocr/v1/general" # 标准精度带位置
|
|
headers = {
|
|
'Content-Type': 'application/json'
|
|
}
|
|
params = {
|
|
'access_token': get_access_token(),
|
|
'image': image_base64,
|
|
'probability': 'true'
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, data=params, timeout=5)
|
|
|
|
try:
|
|
return response.json()
|
|
except requests.exceptions.JSONDecodeError:
|
|
return None
|
|
|
|
except requests.exceptions.Timeout:
|
|
return None
|
|
except requests.exceptions.RequestException as e:
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
logger.info("ocr server 开始加载")
|
|
cfg = toml.load('/home/evan/Workplace/project_infer/cfg_infer_server.toml')
|
|
|
|
# 配置日志输出
|
|
logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention = 5, level="INFO")
|
|
|
|
# context1 = zmq.Context()
|
|
# camera_socket = context1.socket(zmq.REQ)
|
|
# camera_socket.connect(f"tcp://localhost:{cfg['camera']['camera2_port']}")
|
|
# logger.info("connect camera success")
|
|
|
|
# cap = cv2.VideoCapture(20)
|
|
# cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
|
|
# cap.set(cv2.CAP_PROP_FPS, 20)
|
|
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960)
|
|
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 540)
|
|
|
|
# 初始化 server
|
|
context = zmq.Context()
|
|
# 启动 server
|
|
socket = context.socket(zmq.REP)
|
|
socket.bind(f"tcp://*:{cfg['server']['ocr_infer_port']}")
|
|
|
|
import signal
|
|
import sys
|
|
def signal_handler(signum, frame):
|
|
logger.info(f"接收到退出信号 {signum}, 退出中")
|
|
socket.close()
|
|
context.term()
|
|
sys.exit(0)
|
|
|
|
# 注册信号处理函数
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
|
|
logger.info("ocr server 加载完成")
|
|
|
|
while True:
|
|
message1 = socket.recv_string()
|
|
logger.info("收到客户端请求")
|
|
|
|
logger.info("构造摄像头")
|
|
cap = cv2.VideoCapture(20)
|
|
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M','J','P','G'))
|
|
cap.set(cv2.CAP_PROP_FPS, 20)
|
|
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960)
|
|
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 540)
|
|
|
|
for _ in range(5):
|
|
ret, frame = cap.read()
|
|
cv2.waitKey(1)
|
|
|
|
if ret:
|
|
frame = frame[:,0:480]
|
|
frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
|
|
|
now = datetime.datetime.now()
|
|
timestamp = now.strftime("%Y%m%d_%H%M%S")
|
|
output_file_path = f"./saved_picture/{timestamp}.jpg"
|
|
success = cv2.imwrite(output_file_path, frame)
|
|
|
|
_, frame = cv2.imencode('.jpg', frame)
|
|
|
|
encoded_image = base64.b64encode(frame).decode('utf-8')
|
|
|
|
result = ocr_api_request(encoded_image)
|
|
# print(result)
|
|
if result != None:
|
|
socket.send_pyobj({'code': 0, 'content': result.get('words_result')})
|
|
logger.info(f"ocr 返回 {result.get('words_result')}")
|
|
else:
|
|
socket.send_pyobj({'code': -1, 'content': " ocr 没找到文字"})
|
|
logger.error("ocr 没找到文字")
|
|
else:
|
|
socket.send_pyobj({'code': -1, 'content': "ocr 摄像头读取出错"})
|
|
logger.critical("ocr 摄像头读取出错")
|
|
|
|
cap.release()
|
|
|
|
if cv2.waitKey(1) == 27:
|
|
break
|
|
logger.info("ocr infer server exit")
|
|
|