(cfg_det, cfg_rec)
| 123 | |
| 124 | |
| 125 | def main(cfg_det, cfg_rec): |
| 126 | img_path = './testA/' |
| 127 | image_file_list = get_image_file_list(img_path) |
| 128 | drop_score = 0.5 |
| 129 | text_sys = OpenOCRParallel( |
| 130 | drop_score=drop_score, |
| 131 | det_box_type='quad') # det_box_type: 'quad' or 'poly' |
| 132 | is_visualize = False |
| 133 | if is_visualize: |
| 134 | font_path = './simfang.ttf' |
| 135 | check_and_download_font(font_path) |
| 136 | draw_img_save_dir = img_path + 'e2e_results/' if img_path[ |
| 137 | -1] != '/' else img_path[:-1] + 'e2e_results/' |
| 138 | os.makedirs(draw_img_save_dir, exist_ok=True) |
| 139 | save_results = [] |
| 140 | |
| 141 | # Prepare images |
| 142 | images = [] |
| 143 | t_start = time.time() |
| 144 | for image_file in image_file_list: |
| 145 | img, flag_gif, flag_pdf = check_and_read(image_file) |
| 146 | if not flag_gif and not flag_pdf: |
| 147 | img = cv2.imread(image_file) |
| 148 | if img is not None: |
| 149 | images.append((img, img.copy())) |
| 150 | |
| 151 | results = text_sys.process_images(images) |
| 152 | print(f'time cost: {time.time() - t_start}') |
| 153 | # Save results and visualize |
| 154 | for image_id, res in results.items(): |
| 155 | image_file = image_file_list[image_id] |
| 156 | save_pred = f'{os.path.basename(image_file)}\t{json.dumps(res, ensure_ascii=False)}\n' |
| 157 | # print(save_pred) |
| 158 | save_results.append(save_pred) |
| 159 | |
| 160 | if is_visualize: |
| 161 | dt_boxes = [result['points'] for result in res] |
| 162 | rec_res = [result['transcription'] for result in res] |
| 163 | rec_score = [result['score'] for result in res] |
| 164 | image = Image.fromarray( |
| 165 | cv2.cvtColor(images[image_id][0], cv2.COLOR_BGR2RGB)) |
| 166 | draw_img = draw_ocr_box_txt(image, |
| 167 | dt_boxes, |
| 168 | rec_res, |
| 169 | rec_score, |
| 170 | drop_score=drop_score, |
| 171 | font_path=font_path) |
| 172 | |
| 173 | save_file = os.path.join(draw_img_save_dir, |
| 174 | os.path.basename(image_file)) |
| 175 | cv2.imwrite(save_file, draw_img[:, :, ::-1]) |
| 176 | |
| 177 | with open(os.path.join(draw_img_save_dir, 'system_results.txt'), |
| 178 | 'w', |
| 179 | encoding='utf-8') as f: |
| 180 | f.writelines(save_results) |
| 181 | |
| 182 |
no test coverage detected