|
|
|
@@ -0,0 +1,140 @@
|
|
|
|
|
# -- encoding: utf-8 --
|
|
|
|
|
|
|
|
|
|
"""
|
|
|
|
|
Description:
|
|
|
|
|
Create: 2024/1/30 9:26
|
|
|
|
|
"""
|
|
|
|
|
import math
|
|
|
|
|
import time
|
|
|
|
|
from typing import Dict, Any
|
|
|
|
|
|
|
|
|
|
import cv2
|
|
|
|
|
import numpy as np
|
|
|
|
|
from loguru import logger
|
|
|
|
|
|
|
|
|
|
from deepdata.common.utils import bytes_transform
|
|
|
|
|
from deepdata.core.base_op import Mapper
|
|
|
|
|
|
|
|
|
|
from .base_model import BaseModel
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ImgDirectionCorrect(Mapper):
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
super(ImgDirectionCorrect, self).__init__(*args, **kwargs)
|
|
|
|
|
self.img_resize = 1000
|
|
|
|
|
self.limit_size = 30000
|
|
|
|
|
self.use_model = True
|
|
|
|
|
self.model = self.get_model(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _detect_angle(img):
|
|
|
|
|
"""检测图片倾斜角度"""
|
|
|
|
|
# 转为灰度单通道 [[255 255],[255 255]]
|
|
|
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
|
# 黑白颠倒
|
|
|
|
|
gray = cv2.bitwise_not(gray)
|
|
|
|
|
# 二值化
|
|
|
|
|
thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
|
|
|
|
|
# 把大于0的点的行列找出来
|
|
|
|
|
ys, xs = np.where(thresh > 0)
|
|
|
|
|
# 组成坐标[[306 37][306 38][307 38]],里面都是非零的像素
|
|
|
|
|
coords = np.column_stack([xs, ys])
|
|
|
|
|
# 获取最小矩形的信息 返回值(中心点，长宽，角度)
|
|
|
|
|
rect = cv2.minAreaRect(coords)
|
|
|
|
|
# 这里minAreaRect返回值为【0,90】，离y轴最近的夹角，后续有优化空间
|
|
|
|
|
# 夹角小于45度时，填充的空白较少，有助于提升识别率
|
|
|
|
|
angle = rect[-1] # 最后一个参数是角度
|
|
|
|
|
# 小于45度时，逆时针旋转45度
|
|
|
|
|
if angle <= 45.0:
|
|
|
|
|
return angle
|
|
|
|
|
# 大于45度时，顺时针旋转（90-angle）
|
|
|
|
|
return angle - 90
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _detect_direction(image, file_name, model):
|
|
|
|
|
"""
|
|
|
|
|
Args:
|
|
|
|
|
image: 待预测的图片
|
|
|
|
|
file_name: 文件名
|
|
|
|
|
model: 使用的模型， vertical_model 和 standard_model
|
|
|
|
|
Returns: 旋转后的图片
|
|
|
|
|
"""
|
|
|
|
|
# cls_res为模型预测结果，格式应当类似于: [('90', 0.9815167)]
|
|
|
|
|
cls_res = model.infer.predict([image])[0]
|
|
|
|
|
rotate_angle = int(cls_res.get("class_ids", np.array([0], dtype='int32')).item())
|
|
|
|
|
pro = float(cls_res.get("scores", np.array([0], dtype='int32')).item())
|
|
|
|
|
logger.info(
|
|
|
|
|
f"fileName: {file_name}, model detect result is {rotate_angle} with confidence {pro}")
|
|
|
|
|
if rotate_angle == 90 and pro > 0.89:
|
|
|
|
|
return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
|
|
|
|
|
if rotate_angle == 180 and pro > 0.89:
|
|
|
|
|
return cv2.rotate(image, cv2.ROTATE_180)
|
|
|
|
|
if rotate_angle == 270 and pro > 0.89:
|
|
|
|
|
return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
|
|
|
|
return image
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _rotate_bound(image, angle):
|
|
|
|
|
"""根据倾斜角度旋转图片
|
|
|
|
|
Args:
|
|
|
|
|
image: 待处理图片
|
|
|
|
|
angle: _detect_angle方法检测到的倾斜角
|
|
|
|
|
"""
|
|
|
|
|
if angle == 0.0:
|
|
|
|
|
return image
|
|
|
|
|
# 获取宽高
|
|
|
|
|
h, w = image.shape[:2]
|
|
|
|
|
sinval = math.fabs(math.sin(angle))
|
|
|
|
|
cosval = math.fabs(math.cos(angle))
|
|
|
|
|
dx = max(int((w * cosval + h * sinval - w) / 2), 0)
|
|
|
|
|
dy = max(int((w * sinval + h * cosval - h) / 2), 0)
|
|
|
|
|
dst_img = cv2.copyMakeBorder(image, dy, dy, dx, dx, cv2.BORDER_CONSTANT, value=(255, 255, 255))
|
|
|
|
|
h, w = dst_img.shape[:2]
|
|
|
|
|
rotated_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0)
|
|
|
|
|
dst_img = cv2.warpAffine(dst_img, rotated_matrix, (w, h), borderValue=(255, 255, 255))
|
|
|
|
|
return dst_img
|
|
|
|
|
|
|
|
|
|
def init_model(self, *args, **kwargs):
|
|
|
|
|
return BaseModel(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
def execute(self, sample: Dict[str, Any]):
|
|
|
|
|
start = time.time()
|
|
|
|
|
self.read_file_first(sample)
|
|
|
|
|
file_name = sample[self.filename_key]
|
|
|
|
|
file_type = "." + sample[self.filetype_key]
|
|
|
|
|
img_bytes = sample[self.data_key]
|
|
|
|
|
if img_bytes:
|
|
|
|
|
data = bytes_transform.bytes_to_numpy(img_bytes)
|
|
|
|
|
correct_data = self._img_direction_correct(data, file_name, self.model)
|
|
|
|
|
sample[self.data_key] = bytes_transform.numpy_to_bytes(correct_data, file_type)
|
|
|
|
|
logger.info(f"fileName: {file_name}, method: ImgDirectionCorrect costs {time.time() - start:6f} s")
|
|
|
|
|
return sample
|
|
|
|
|
|
|
|
|
|
def _img_direction_correct(self, img, file_name, standard_model):
|
|
|
|
|
height, width = img.shape[:2]
|
|
|
|
|
if max(height, width) > self.limit_size:
|
|
|
|
|
logger.info(
|
|
|
|
|
f"fileName: {file_name}, method: ImgDirectionCorrect cannot process pixels number larger than 30000")
|
|
|
|
|
return img
|
|
|
|
|
detect_angle_img = self._resize(img)
|
|
|
|
|
# 检测旋转角
|
|
|
|
|
angle = self._detect_angle(detect_angle_img)
|
|
|
|
|
# 将图片处理为 0, 90, 180, 270旋转角度的图片
|
|
|
|
|
rotated_img = self._rotate_bound(img, angle)
|
|
|
|
|
# 0-180方向识别：二分类模型，检测图片方向角为 0, 180, 将其处理为 0和180二分类图片
|
|
|
|
|
rotated_img = self._detect_direction(rotated_img, file_name, standard_model)
|
|
|
|
|
return rotated_img
|
|
|
|
|
|
|
|
|
|
def _resize(self, image):
|
|
|
|
|
height, width = image.shape[:2] # 获取原图像的水平方向尺寸和垂直方向尺寸。
|
|
|
|
|
temp = max(height, width)
|
|
|
|
|
# 若图片最长边大于限值，对图片进行压缩，否则返回原图
|
|
|
|
|
if temp >= self.img_resize:
|
|
|
|
|
mul_temp = temp / self.img_resize
|
|
|
|
|
if height > width:
|
|
|
|
|
return cv2.resize(image, (int(width / mul_temp), self.img_resize), interpolation=cv2.INTER_AREA)
|
|
|
|
|
elif height < width:
|
|
|
|
|
return cv2.resize(image, (self.img_resize, int(height / mul_temp)), interpolation=cv2.INTER_AREA)
|
|
|
|
|
else:
|
|
|
|
|
return cv2.resize(image, (self.img_resize, self.img_resize), interpolation=cv2.INTER_AREA)
|
|
|
|
|
return image
|
figmar