优化照片涂抹功能架构

util/util.py

@@ -1,6 +1,78 @@
+import logging
+import os
 from datetime import datetime
 
 
 # 获取yyyy-MM-dd HH:mm:ss格式的当前时间
 def get_default_datetime():
     return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+
+
+def get_ocr_layout(ocr, img_path):
+    """
+    获取ocr识别的结果，转为合适的layout形式
+    :param ocr: ocr模型
+    :param img_path: 图片本地路径
+    :return:
+    """
+
+    def _get_box(old_box):
+        new_box = [
+            min(old_box[0][0], old_box[3][0]),  # x1
+            min(old_box[0][1], old_box[1][1]),  # y1
+            max(old_box[1][0], old_box[2][0]),  # x2
+            max(old_box[2][1], old_box[3][1]),  # y2
+        ]
+        return new_box
+
+    def _normal_box(box_data):
+        # Ensure the height and width of bbox are greater than zero
+        if box_data[3] - box_data[1] < 0 or box_data[2] - box_data[0] < 0:
+            return False
+        return True
+
+    layout = []
+    ocr_result = ocr.ocr(img_path, cls=False)
+    ocr_result = ocr_result[0]
+    if not ocr_result:
+        return layout
+    for segment in ocr_result:
+        box = segment[0]
+        box = _get_box(box)
+        if not _normal_box(box):
+            continue
+        text = segment[1][0]
+        layout.append((box, text))
+    return layout
+
+
+def delete_temp_file(temp_files):
+    """
+    删除临时文件，可以批量
+    :param temp_files: 临时文件路径
+    """
+    if not temp_files:
+        return
+    if isinstance(temp_files, str):
+        temp_files = [temp_files]
+    for file in temp_files:
+        try:
+            os.remove(file)
+            logging.info(f"临时文件 {file} 已删除")
+        except Exception as e:
+            logging.warning(f"删除临时文件 {file} 时出错: {e}")
+
+
+def zoom_rectangle(rectangle, ratio):
+    """
+    缩放矩形
+    :param rectangle: 原矩形坐标
+    :param ratio: 缩放比率
+    :return: 缩放后的矩形坐标
+    """
+    x1, y1, x2, y2 = rectangle
+    x1 = round(x1 - x1 * ratio)
+    y1 = round(y1 - y1 * ratio)
+    x2 = round(x2 + x2 * ratio)
+    y2 = round(y2 + y2 * ratio)
+    return [x1, y1, x2, y2]