| 25 | |
| 26 | |
| 27 | class OpenOCRParallel: |
| 28 | |
| 29 | def __init__(self, drop_score=0.5, det_box_type='quad', max_rec_threads=1): |
| 30 | cfg_det = Config( |
| 31 | './configs/det/dbnet/repvit_db.yml').cfg # mobile model |
| 32 | # cfg_rec = Config('./configs/rec/svtrv2/svtrv2_ch.yml').cfg # server model |
| 33 | cfg_rec = Config( |
| 34 | './configs/rec/svtrv2/repsvtr_ch.yml').cfg # mobile model |
| 35 | self.text_detector = OpenDetector(cfg_det, numId=0) |
| 36 | self.text_recognizer = OpenRecognizer(cfg_rec, numId=0) |
| 37 | self.det_box_type = det_box_type |
| 38 | self.drop_score = drop_score |
| 39 | self.queue = queue.Queue( |
| 40 | ) # Queue to hold detected boxes for recognition |
| 41 | self.results = {} |
| 42 | self.lock = threading.Lock() # Lock for thread-safe access to results |
| 43 | self.max_rec_threads = max_rec_threads |
| 44 | self.stop_signal = threading.Event() # Signal to stop threads |
| 45 | |
| 46 | def start_recognition_threads(self): |
| 47 | """Start recognition threads.""" |
| 48 | self.rec_threads = [] |
| 49 | for _ in range(self.max_rec_threads): |
| 50 | t = threading.Thread(target=self.recognize_text) |
| 51 | t.start() |
| 52 | self.rec_threads.append(t) |
| 53 | |
| 54 | def detect_text(self, image_list): |
| 55 | """Single-threaded text detection for all images.""" |
| 56 | for image_id, (img_numpy, ori_img) in enumerate(image_list): |
| 57 | dt_boxes = self.text_detector(img_numpy=img_numpy)[0]['boxes'] |
| 58 | if dt_boxes is None: |
| 59 | self.results[image_id] = [] # If no boxes, set empty results |
| 60 | continue |
| 61 | |
| 62 | dt_boxes = sorted_boxes(dt_boxes) |
| 63 | img_crop_list = [] |
| 64 | for box in dt_boxes: |
| 65 | tmp_box = np.array(box).astype(np.float32) |
| 66 | img_crop = (get_rotate_crop_image(ori_img, tmp_box) |
| 67 | if self.det_box_type == 'quad' else |
| 68 | get_minarea_rect_crop(ori_img, tmp_box)) |
| 69 | img_crop_list.append(img_crop) |
| 70 | self.queue.put( |
| 71 | (image_id, dt_boxes, img_crop_list |
| 72 | )) # Put image ID, detected box, and cropped image in queue |
| 73 | |
| 74 | # Signal that no more items will be added to the queue |
| 75 | self.stop_signal.set() |
| 76 | |
| 77 | def recognize_text(self): |
| 78 | """Recognize text in each cropped image.""" |
| 79 | while not self.stop_signal.is_set() or not self.queue.empty(): |
| 80 | try: |
| 81 | image_id, boxs, img_crop_list = self.queue.get(timeout=0.5) |
| 82 | rec_results = self.text_recognizer( |
| 83 | img_numpy_list=img_crop_list, batch_num=6) |
| 84 | for rec_result, box in zip(rec_results, boxs): |