调整ocr调用

This commit is contained in:
2024-06-24 14:26:03 +08:00
parent 6e77f20c1e
commit 88666ccf88

View File

@@ -46,20 +46,18 @@ def visual_model_test(model_type, test_img, task_path, schema):
imgs = split_image(test_img) imgs = split_image(test_img)
layout = [] layout = []
temp_files_paths = [] temp_files_paths = []
doc_parser = DocParser(layout_analysis=False)
for img in imgs: for img in imgs:
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
cv2.imwrite(temp_file.name, img["img"]) cv2.imwrite(temp_file.name, img["img"])
temp_files_paths.append(temp_file.name) temp_files_paths.append(temp_file.name)
parsed_doc = doc_parser.parse({"doc": temp_file.name}, expand_to_a4_size=True) ocr_layout = get_ocr_layout(temp_file.name)
if img["x_offset"] or img["y_offset"]: if img["x_offset"] or img["y_offset"]:
for p in parsed_doc["layout"]: for box in ocr_layout:
box = p[0]
box[0] += img["x_offset"] box[0] += img["x_offset"]
box[1] += img["y_offset"] box[1] += img["y_offset"]
box[2] += img["x_offset"] box[2] += img["x_offset"]
box[3] += img["y_offset"] box[3] += img["y_offset"]
layout += parsed_doc["layout"] layout += ocr_layout
write_visual_result(test_img, layout=layout) write_visual_result(test_img, layout=layout)
else: else: