import base64
import json
import os
import cv2
import numpy as np
import yaml
from mosec import Worker, Server
from mosec.errors import ValidationError
from inference import Inference
def _get_model():
if os.path.exists('config.yaml'):
config = open('config.yaml', mode='r', encoding='utf-8')
config = yaml.load(config, Loader=yaml.FullLoader)
model = Inference(
det_model_path=config['det_model_path'],
rec_model_path=config['rec_model_path'],
device=config['device'],
dict_path=config['dict_path'],
rec_std=0.5, rec_mean=0.5, threshold=0.7,
angle_classes=config['angle_classes'],
angle_classify_model_path=config['angle_model_path'],
object_classes=None,
object_classify_model_path=None
)
return model, config
else:
raise FileNotFoundError('must have a config.yaml file!')
class OCRInference(Worker):
def __init__(self):
super(OCRInference, self).__init__()
self.model, self.config = _get_model()
def forward(self, req: dict):
try:
image = req["image"]
save_name = req['saveName']
im = np.frombuffer(base64.b64decode(image), np.uint8)
im = cv2.imdecode(im, 1)
result = self.model.infer(img=im,
img_save_name=save_name,
cut_image_save_path=self.config['cut_image_save_path'],
need_angle=self.config['need_angle'],
need_object=self.config['need_object'])
return json.dumps({'status': 1, 'result': result})
except KeyError as err:
raise ValidationError(f"cannot find key {err}")
except Exception as err:
raise ValidationError(f"cannot decode as image data: {err}")
if __name__ == "__main__":
server = Server()
server.append_worker(OCRInference, num=2, max_batch_size=16)
server.run()