35 lines
767 B
Python
35 lines
767 B
Python
|
|
||
|
|
from keypoint_detector import RTMPose
|
||
|
|
from rapidocr_openvino import RapidOCR
|
||
|
|
|
||
|
|
from utils import four_point_transform
|
||
|
|
from result import Result
|
||
|
|
from configs.config_certificate import config_certificate
|
||
|
|
|
||
|
|
|
||
|
|
class PipePredictor():
|
||
|
|
def __init__(self) -> None:
|
||
|
|
|
||
|
|
self.kpt_detector = RTMPose('models/kpt_certificate_encrypt')
|
||
|
|
self.ocr = RapidOCR(config_certificate)
|
||
|
|
|
||
|
|
self.result = Result()
|
||
|
|
|
||
|
|
def __call__(self, img):
|
||
|
|
self.result.clear()
|
||
|
|
keypoints, scores = self.kpt_detector(img)
|
||
|
|
|
||
|
|
self.warped_img = four_point_transform(img, keypoints[0])
|
||
|
|
|
||
|
|
self.ocr_res, _ = self.ocr(self.warped_img)
|
||
|
|
# print(self.ocr_res)
|
||
|
|
self.result.update(self.ocr_res)
|
||
|
|
return self.result.get()
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|