(img_path)
| 48 | cv2.imwrite(save_path, img) |
| 49 | |
| 50 | def ocr(img_path): |
| 51 | while True: |
| 52 | try: |
| 53 | vlm_output = vlm.generate(query=prompt, image=[img_path]) |
| 54 | vlm_output = vlm_output.replace('```json', '') |
| 55 | vlm_output = vlm_output.replace('```', '') |
| 56 | boxes = [] |
| 57 | vlm_output = json.loads(vlm_output) |
| 58 | for obj in vlm_output['objects']: |
| 59 | if obj['type'] == 'object': |
| 60 | continue |
| 61 | if 'content' not in obj: |
| 62 | raise Exception('not good') |
| 63 | box = obj['bounding_box'] |
| 64 | boxes.append(box) |
| 65 | if any([any([x < 0 or x > 1000 for x in box]) for box in boxes]): |
| 66 | raise Exception('beyond the image') |
| 67 | draw_boxes(img_path, boxes) |
| 68 | # save json |
| 69 | with open(img_path.replace('img', 'vlmocr').replace('.jpg', '.json'), 'w') as f: |
| 70 | json.dump(vlm_output, f, indent=2, ensure_ascii=False) |
| 71 | break |
| 72 | except Exception as e: |
| 73 | print(f'Error in {img_path}, retrying...') |
| 74 | time.sleep(2) |
| 75 | |
| 76 | |
| 77 |
no test coverage detected