diff --git a/check_ie_result/check_ie_result.py b/check_ie_result/check_ie_result.py index 83388ef..7125387 100644 --- a/check_ie_result/check_ie_result.py +++ b/check_ie_result/check_ie_result.py @@ -14,7 +14,7 @@ from ucloud import ufile from util import image_util -def check_ie_result(pk_phhd): +def check_ie_result(pk_phhd, need_to_annotation=True): os.makedirs(f"./check_result/{pk_phhd}", exist_ok=True) json_result = {"pk_phhd": pk_phhd} session = MysqlSession() @@ -46,45 +46,51 @@ def check_ie_result(pk_phhd): ZxPhrec.pk_phhd == pk_phhd).all() for phrec in phrecs: img_name = phrec.cfjaddress - img_path = ufile.get_private_url(img_name) + img_path = ufile.get_private_url(img_name, "drg2015") + if not img_path: + img_path = ufile.get_private_url(img_name) response = requests.get(img_path) image = Image.open(BytesIO(response.content)).convert("RGB") - font_size = image.width * image.height / 200000 - font = ImageFont.truetype("./font/simfang.ttf", size=font_size) + if need_to_annotation: + font_size = image.width * image.height / 200000 + font = ImageFont.truetype("./font/simfang.ttf", size=font_size) - ocr = session.query(ZxIeResult.id, ZxIeResult.content, ZxIeResult.rotation_angle, ZxIeResult.x_offset, - ZxIeResult.y_offset).filter(ZxIeResult.pk_phrec == phrec.pk_phrec).all() - if not ocr: + ocr = session.query(ZxIeResult.id, ZxIeResult.content, ZxIeResult.rotation_angle, ZxIeResult.x_offset, + ZxIeResult.y_offset).filter(ZxIeResult.pk_phrec == phrec.pk_phrec).all() + if not ocr: + os.makedirs(f"./check_result/{pk_phhd}/0", exist_ok=True) + image.save(f"./check_result/{pk_phhd}/0/{img_name}") + + for _, group_results in groupby(ocr, key=lambda x: x.id): + draw = ImageDraw.Draw(image) + for ocr_item in group_results: + result = json.loads(ocr_item.content) + rotation_angle = ocr_item.rotation_angle + x_offset = ocr_item.x_offset + y_offset = ocr_item.y_offset + for key in result: + for value in result[key]: + box = value["bbox"][0] + + if rotation_angle: + box = image_util.invert_rotate_rectangle(box, (image.width / 2, image.height / 2), + rotation_angle) + if x_offset: + box[0] += x_offset + box[2] += x_offset + if y_offset: + box[1] += y_offset + box[3] += y_offset + + draw.rectangle(box, outline="red", width=2) # 绘制矩形 + draw.text((box[0], box[1] - font_size), key, fill="blue", font=font) # 在矩形上方绘制文本 + draw.text((box[0], box[3]), value["text"], fill="blue", font=font) # 在矩形下方绘制文本 + os.makedirs(f"./check_result/{pk_phhd}/{ocr_item.id}", exist_ok=True) + image.save(f"./check_result/{pk_phhd}/{ocr_item.id}/{img_name}") + else: os.makedirs(f"./check_result/{pk_phhd}/0", exist_ok=True) image.save(f"./check_result/{pk_phhd}/0/{img_name}") - - for _, group_results in groupby(ocr, key=lambda x: x.id): - draw = ImageDraw.Draw(image) - for ocr_item in group_results: - result = json.loads(ocr_item.content) - rotation_angle = ocr_item.rotation_angle - x_offset = ocr_item.x_offset - y_offset = ocr_item.y_offset - for key in result: - for value in result[key]: - box = value["bbox"][0] - - if rotation_angle: - box = image_util.invert_rotate_rectangle(box, (image.width / 2, image.height / 2), - rotation_angle) - if x_offset: - box[0] += x_offset - box[2] += x_offset - if y_offset: - box[1] += y_offset - box[3] += y_offset - - draw.rectangle(box, outline="red", width=2) # 绘制矩形 - draw.text((box[0], box[1] - font_size), key, fill="blue", font=font) # 在矩形上方绘制文本 - draw.text((box[0], box[3]), value["text"], fill="blue", font=font) # 在矩形下方绘制文本 - os.makedirs(f"./check_result/{pk_phhd}/{ocr_item.id}", exist_ok=True) - image.save(f"./check_result/{pk_phhd}/{ocr_item.id}/{img_name}") session.close() # 自定义JSON处理器 @@ -99,4 +105,4 @@ def check_ie_result(pk_phhd): if __name__ == '__main__': - check_ie_result(0) + check_ie_result(5640504)