33# @Contact: liekkaskono@163.com
44import argparse
55import base64
6- import importlib
76import io
87import os
98import sys
1413import uvicorn
1514from fastapi import FastAPI , Form , UploadFile
1615from PIL import Image
17-
18- if importlib .util .find_spec ("rapidocr_onnxruntime" ):
19- from rapidocr_onnxruntime import RapidOCR
20- elif importlib .util .find_spec ("rapidocr_paddle" ):
21- from rapidocr_paddle import RapidOCR
22- elif importlib .util .find_spec ("rapidocr_openvino" ):
23- from rapidocr_openvino import RapidOCR
24- else :
25- raise ImportError (
26- "Please install one of [rapidocr_onnxruntime,rapidocr-paddle,rapidocr-openvino]"
27- )
16+ from rapidocr import RapidOCR
2817
2918sys .path .append (str (Path (__file__ ).resolve ().parent .parent ))
3019
@@ -39,36 +28,29 @@ def __init__(self) -> None:
3928 self .ocr = RapidOCR ()
4029 else :
4130 self .ocr = RapidOCR (
42- det_model_path = det_model_path ,
43- cls_model_path = cls_model_path ,
44- rec_model_path = rec_model_path ,
31+ params = {
32+ "Det.model_path" : det_model_path ,
33+ "Cls.model_path" : cls_model_path ,
34+ "Rec.model_path" : rec_model_path ,
35+ }
4536 )
4637
4738 def __call__ (
48- self , img : Image .Image , use_det = None , use_cls = None , use_rec = None , ** kwargs
39+ self , ori_img : Image .Image , use_det = None , use_cls = None , use_rec = None , ** kwargs
4940 ) -> Dict :
50- img = np .array (img )
51- ocr_res , _ = self .ocr (
41+ img = np .array (ori_img )
42+ ocr_res = self .ocr (
5243 img , use_det = use_det , use_cls = use_cls , use_rec = use_rec , ** kwargs
5344 )
5445
55- if not ocr_res :
46+ if ocr_res . boxes is None or ocr_res . txts is None or ocr_res . scores is None :
5647 return {}
5748
5849 out_dict = {}
59- for i , dats in enumerate (ocr_res ):
60- values = {}
61- for dat in dats :
62- if isinstance (dat , str ):
63- values ["rec_txt" ] = dat
64-
65- if isinstance (dat , np .float32 ):
66- values ["score" ] = f"{ dat :.4f} "
67-
68- if isinstance (dat , list ):
69- values ["dt_boxes" ] = dat
70- out_dict [str (i )] = values
71-
50+ for i , (boxes , txt , score ) in enumerate (
51+ zip (ocr_res .boxes , ocr_res .txts , ocr_res .scores )
52+ ):
53+ out_dict [i ] = {"rec_txt" : txt , "dt_boxes" : boxes .tolist (), "score" : score }
7254 return out_dict
7355
7456
@@ -100,6 +82,7 @@ def ocr(
10082 "When sending a post request, data or files must have a value."
10183 )
10284 ocr_res = processor (img , use_det = use_det , use_cls = use_cls , use_rec = use_rec )
85+
10386 return ocr_res
10487
10588
0 commit comments