(input_image,
model_type_select,
det_input_size_textbox=960,
rec_drop_score=0.4,
mask_thresh=0.3,
box_thresh=0.6,
unclip_ratio=1.5,
det_score_mode='slow')
| 35 | |
| 36 | |
| 37 | def main(input_image, |
| 38 | model_type_select, |
| 39 | det_input_size_textbox=960, |
| 40 | rec_drop_score=0.4, |
| 41 | mask_thresh=0.3, |
| 42 | box_thresh=0.6, |
| 43 | unclip_ratio=1.5, |
| 44 | det_score_mode='slow'): |
| 45 | global text_sys, model_type |
| 46 | |
| 47 | # Update OCR model if the model type changes |
| 48 | if model_type_select != model_type: |
| 49 | model_type = model_type_select |
| 50 | text_sys = initialize_ocr(model_type, rec_drop_score) |
| 51 | |
| 52 | img = input_image[:, :, ::-1] |
| 53 | starttime = time.time() |
| 54 | results, time_dict, mask = text_sys( |
| 55 | img_numpy=img, |
| 56 | return_mask=True, |
| 57 | det_input_size=int(det_input_size_textbox), |
| 58 | thresh=mask_thresh, |
| 59 | box_thresh=box_thresh, |
| 60 | unclip_ratio=unclip_ratio, |
| 61 | score_mode=det_score_mode) |
| 62 | elapse = time.time() - starttime |
| 63 | save_pred = json.dumps(results[0], ensure_ascii=False) |
| 64 | image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) |
| 65 | boxes = [res['points'] for res in results[0]] |
| 66 | txts = [res['transcription'] for res in results[0]] |
| 67 | scores = [res['score'] for res in results[0]] |
| 68 | draw_img = draw_ocr_box_txt( |
| 69 | image, |
| 70 | boxes, |
| 71 | txts, |
| 72 | scores, |
| 73 | drop_score=rec_drop_score, |
| 74 | font_path=font_path, |
| 75 | ) |
| 76 | mask = mask[0, 0, :, :] > mask_thresh |
| 77 | return save_pred, elapse, draw_img, mask.astype('uint8') * 255 |
| 78 | |
| 79 | |
| 80 | def get_all_file_names_including_subdirs(dir_path): |
nothing calls this directly
no test coverage detected