281 lines
9.2 KiB
Python
281 lines
9.2 KiB
Python
import logging
|
|
import math
|
|
import os
|
|
|
|
import cv2
|
|
import numpy
|
|
import requests
|
|
from paddleclas import PaddleClas
|
|
from tenacity import retry, stop_after_attempt, wait_random
|
|
|
|
from log import PROJECT_ROOT
|
|
|
|
|
|
def capture(image, rectangle):
|
|
"""
|
|
截取图片
|
|
:param image: ndarray
|
|
:param rectangle: 要截取的矩形
|
|
:return: 截取之后的ndarray图片
|
|
"""
|
|
x1, y1, x2, y2 = rectangle
|
|
height, width = image.shape[:2]
|
|
# 确保坐标值在图片范围内
|
|
x1 = max(0, x1)
|
|
y1 = max(0, y1)
|
|
x2 = min(width, x2)
|
|
y2 = min(height, y2)
|
|
return image[int(y1):int(y2), int(x1):int(x2)]
|
|
|
|
|
|
def split(img_path, ratio=1.414, overlap=0.05, x_compensation=3):
|
|
"""
|
|
分割图片
|
|
:param img_path:图片路径
|
|
:param ratio: 分割后的比例
|
|
:param overlap: 图片之间的覆盖比例
|
|
:param x_compensation: 横向补偿倍率
|
|
:return: 分割后的图片组(NumPy数组形式)
|
|
"""
|
|
split_result = []
|
|
image = cv2.imread(img_path)
|
|
height, width = image.shape[:2]
|
|
hw_ratio = height / width
|
|
wh_ratio = width / height
|
|
|
|
img_name, img_ext = parse_save_path(img_path)
|
|
if hw_ratio > ratio: # 纵向过长
|
|
new_img_height = width * ratio
|
|
step = width * (ratio - overlap) # 偏移步长
|
|
for i in range(math.ceil(height / step)):
|
|
offset = round(step * i)
|
|
cropped_img = capture(image, [0, offset, width, offset + new_img_height])
|
|
split_path = get_save_path(f'{img_name}.split_{i}.{img_ext}')
|
|
cv2.imwrite(split_path, cropped_img)
|
|
split_result.append({'img': split_path, 'x_offset': 0, 'y_offset': offset})
|
|
elif wh_ratio > ratio: # 横向过长
|
|
new_img_width = height * ratio
|
|
step = height * (ratio - overlap * x_compensation) # 一般文字是横向的，所以横向截取时增大重叠部分
|
|
for i in range(math.ceil(width / step)):
|
|
offset = round(step * i)
|
|
cropped_img = capture(image, [offset, 0, offset + new_img_width, width])
|
|
split_path = get_save_path(f'{img_name}.split_{i}.{img_ext}')
|
|
cv2.imwrite(split_path, cropped_img)
|
|
split_result.append({'img': split_path, 'x_offset': offset, 'y_offset': 0})
|
|
else:
|
|
split_result.append({'img': img_path, 'x_offset': 0, 'y_offset': 0})
|
|
return split_result
|
|
|
|
|
|
def parse_rotation_angles(image):
|
|
"""
|
|
判断图片旋转角度，逆时针旋转该角度后为正。可能值['0', '90', '180', '270']
|
|
:param image: 图片NumPy数组或文件路径
|
|
:return: 最有可能的两个角度
|
|
"""
|
|
angles = ['0', '90']
|
|
model = PaddleClas(model_name='text_image_orientation')
|
|
clas_result = model.predict(input_data=image)
|
|
try:
|
|
clas_result = next(clas_result)[0]
|
|
if clas_result['scores'][0] < 0.5:
|
|
return angles
|
|
angles = clas_result['label_names']
|
|
except Exception as e:
|
|
logging.error('获取图片旋转角度失败', exc_info=e)
|
|
return angles
|
|
|
|
|
|
def rotate(img_path, angle):
|
|
"""
|
|
旋转图片
|
|
:param img_path: 图片NumPy数组
|
|
:param angle: 逆时针旋转角度
|
|
:return: 旋转后的图片NumPy数组
|
|
"""
|
|
if angle == 0:
|
|
return img_path
|
|
image = cv2.imread(img_path)
|
|
height, width = image.shape[:2]
|
|
if angle == 180:
|
|
new_width = width
|
|
new_height = height
|
|
else:
|
|
new_width = height
|
|
new_height = width
|
|
# 绕图像的中心旋转
|
|
# 参数：旋转中心 旋转度数 scale
|
|
matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1)
|
|
# 旋转后平移
|
|
matrix[0, 2] += (new_width - width) / 2
|
|
matrix[1, 2] += (new_height - height) / 2
|
|
# 参数：原始图像 旋转参数 元素图像宽高
|
|
rotated = cv2.warpAffine(image, matrix, (new_width, new_height))
|
|
|
|
img_name, img_ext = parse_save_path(img_path)
|
|
rotated_path = get_save_path(f'{img_name}.rotate_{angle}.{img_ext}')
|
|
cv2.imwrite(rotated_path, rotated)
|
|
return rotated_path
|
|
|
|
|
|
def invert_rotate_point(point, center, angle):
|
|
"""
|
|
反向旋转图片上的点
|
|
:param point: 点
|
|
:param center: 旋转中心
|
|
:param angle: 旋转角度
|
|
:return: 旋转后的点坐标
|
|
"""
|
|
matrix = cv2.getRotationMatrix2D(center, angle, 1)
|
|
if angle != 180:
|
|
# 旋转后平移
|
|
matrix[0, 2] += center[1] - center[0]
|
|
matrix[1, 2] += center[0] - center[1]
|
|
|
|
reverse_matrix = cv2.invertAffineTransform(matrix)
|
|
|
|
point = numpy.array([[point[0]], [point[1]], [1]])
|
|
return numpy.dot(reverse_matrix, point)
|
|
|
|
|
|
def invert_rotate_rectangle(rectangle, center, angle):
|
|
"""
|
|
反向旋转图片上的矩形
|
|
:param rectangle: 矩形
|
|
:param center: 旋转中心
|
|
:param angle: 旋转角度
|
|
:return: 旋转后的矩形坐标
|
|
"""
|
|
if angle == 0:
|
|
return list(rectangle)
|
|
|
|
x1, y1, x2, y2 = rectangle
|
|
|
|
# 计算矩形的四个顶点
|
|
top_left = (x1, y1)
|
|
bot_left = (x1, y2)
|
|
top_right = (x2, y1)
|
|
bot_right = (x2, y2)
|
|
|
|
# 旋转矩形的四个顶点
|
|
rot_top_left = invert_rotate_point(top_left, center, angle).astype(int)
|
|
rot_bot_left = invert_rotate_point(bot_left, center, angle).astype(int)
|
|
rot_bot_right = invert_rotate_point(bot_right, center, angle).astype(int)
|
|
rot_top_right = invert_rotate_point(top_right, center, angle).astype(int)
|
|
|
|
# 找出旋转后矩形的新左上角和右下角坐标
|
|
new_top_left = (min(rot_top_left[0], rot_bot_left[0], rot_bot_right[0], rot_top_right[0]),
|
|
min(rot_top_left[1], rot_bot_left[1], rot_bot_right[1], rot_top_right[1]))
|
|
new_bot_right = (max(rot_top_left[0], rot_bot_left[0], rot_bot_right[0], rot_top_right[0]),
|
|
max(rot_top_left[1], rot_bot_left[1], rot_bot_right[1], rot_top_right[1]))
|
|
|
|
return [new_top_left[0], new_top_left[1], new_bot_right[0], new_bot_right[1]]
|
|
|
|
|
|
def expand_to_a4_size(image):
|
|
"""
|
|
以尽量少的方式将图片扩充到a4大小
|
|
:param image: 图片NumPy数组
|
|
:return: 扩充后的图片NumPy数组和偏移量
|
|
"""
|
|
height, width = image.shape[:2]
|
|
x_offset, y_offset = 0, 0
|
|
hw_ratio = height / width
|
|
if hw_ratio >= 1.42:
|
|
exp_w = int(height / 1.414 - width)
|
|
x_offset = int(exp_w / 2)
|
|
exp_img = numpy.zeros((height, x_offset, 3), dtype='uint8')
|
|
exp_img.fill(255)
|
|
image = numpy.hstack([exp_img, image, exp_img])
|
|
elif 1 <= hw_ratio <= 1.40:
|
|
exp_h = int(width * 1.414 - height)
|
|
y_offset = int(exp_h / 2)
|
|
exp_img = numpy.zeros((y_offset, width, 3), dtype='uint8')
|
|
exp_img.fill(255)
|
|
image = numpy.vstack([exp_img, image, exp_img])
|
|
elif 0.72 <= hw_ratio < 1:
|
|
exp_w = int(height * 1.414 - width)
|
|
x_offset = int(exp_w / 2)
|
|
exp_img = numpy.zeros((height, x_offset, 3), dtype='uint8')
|
|
exp_img.fill(255)
|
|
image = numpy.hstack([exp_img, image, exp_img])
|
|
elif hw_ratio <= 0.7:
|
|
exp_h = int(width / 1.414 - height)
|
|
y_offset = int(exp_h / 2)
|
|
exp_img = numpy.zeros((y_offset, width, 3), dtype='uint8')
|
|
exp_img.fill(255)
|
|
image = numpy.vstack([exp_img, image, exp_img])
|
|
return image, x_offset, y_offset
|
|
|
|
|
|
def combined(img1, img2):
|
|
# 获取两张图片的高度和宽度
|
|
height1, width1 = img1.shape[:2]
|
|
height2, width2 = img2.shape[:2]
|
|
|
|
# 确保两张图片的高度相同
|
|
if height1 != height2:
|
|
# 如果高度不同，调整较小高度的图片
|
|
if height1 < height2:
|
|
img1 = cv2.resize(img1, (int(width1 * height2 / height1), height2))
|
|
else:
|
|
img2 = cv2.resize(img2, (int(width2 * height1 / height2), height1))
|
|
|
|
# 再次获取调整后的图片尺寸
|
|
height1, width1 = img1.shape[:2]
|
|
height2, width2 = img2.shape[:2]
|
|
|
|
# 创建一个空白的图像，宽度等于两张图片的宽度之和，高度等于它们共同的高度
|
|
total_width = width1 + width2
|
|
max_height = max(height1, height2)
|
|
combined_img = numpy.zeros((max_height, total_width, 3), dtype=numpy.uint8)
|
|
|
|
# 将img1和img2复制到新的图像中
|
|
combined_img[:height1, :width1] = img1
|
|
combined_img[:height2, width1:width1 + width2] = img2
|
|
return combined_img
|
|
|
|
|
|
def parse_img_url(url):
|
|
"""
|
|
解析图片url
|
|
:param url: 图片url
|
|
:return: 图片名称和图片后缀
|
|
"""
|
|
url = url.split('?')[0]
|
|
return os.path.basename(url)
|
|
|
|
|
|
@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True,
|
|
after=lambda x: logging.warning('保存图片失败！'))
|
|
def save_to_local(img_url):
|
|
"""
|
|
保存图片到本地
|
|
:param img_url: 图片url
|
|
:return: 本地保存地址
|
|
"""
|
|
response = requests.get(img_url)
|
|
response.raise_for_status() # 检查响应状态码是否正常
|
|
|
|
save_path = get_save_path(parse_img_url(img_url))
|
|
with open(save_path, 'wb') as file:
|
|
file.write(response.content)
|
|
return save_path
|
|
|
|
|
|
def get_img_path(img_full_name):
|
|
save_path = get_save_path(img_full_name)
|
|
if os.path.exists(save_path):
|
|
return save_path
|
|
return None
|
|
|
|
|
|
def get_save_path(img_full_name):
|
|
return os.path.join(PROJECT_ROOT, 'tmp_img', img_full_name)
|
|
|
|
|
|
def parse_save_path(img_path):
|
|
img_full_name = os.path.basename(img_path)
|
|
img_name, img_ext = img_full_name.rsplit('.', 1)
|
|
return img_name, img_ext
|