img-processor/tools/image.py

import os
import numpy as np
import cv2


def is_empty(s: str):
    return s is None or len(str(s)) == 0


def to_int(s: str):
    if isinstance(s, (int, float)):
        return s
    if is_empty(s):
        return None
    return int(str(s))


def to_size_int(s: str):
    v = to_int(s)
    if v is None:
        return v
    if v < 1:
        return None
    return v


class NotImageError(Exception):
    def __init__(self):
        super().__init__("File Is Not Image")


class ProcessorParams:
    def __init__(self, width, height, rotate, flip, rect, out_dir, path):
        self.width = width
        self.height = height
        self.rotate = rotate
        # 翻转， 0 为沿X轴翻转，正数为沿Y轴翻转，负数为同时沿X轴和Y轴翻转
        self.flip = flip
        # 裁剪，left,top,right,bottom 四个参数
        self.rect = rect
        # 保存的文件夹
        self.out_dir = out_dir
        self.path = path

    def set_path(self, path, out_dir):
        self.path = path
        self.out_dir = out_dir



def processor(params: ProcessorParams):
    file_path = params.path

    if os.path.exists(file_path) is not True:
        raise FileNotFoundError('Need to give the source path.')

    if not os.path.exists(params.out_dir):
        os.makedirs(params.out_dir)

    img = cv2.imread(file_path)
    if img is None:
        raise NotImageError()

    origin_h = img.shape[0]
    origin_w = img.shape[1]

    dst, r_w, r_h = crop(img, params.rect, origin_w, origin_h)

    width = to_size_int(params.width)
    height = to_size_int(params.height)
    change_size = True
    if width is None and height is None:
        change_size = False

    if change_size:
        if width is None:
            h = height
            w = round(h * origin_w / origin_h)
        elif height is None:
            w = width
            h = round(w * origin_h / origin_w)
        else:
            w = width
            h = height
        origin = img if dst is None else dst
        dst = cv2.resize(origin, (w, h))
        r_w = w
        r_h = h

    if r_w is None:
        r_w = origin_w
    if r_h is None:
        r_h = origin_h

    flip = params.flip
    if is_empty(flip) is not True:
        flip = to_int(flip)
        origin = img if dst is None else dst
        dst = cv2.flip(origin, flip)

    rotate = params.rotate
    if is_empty(rotate) is False and to_int(rotate) != 0:
        origin = img if dst is None else dst
        dst, new_w, new_h = rot_degree(origin, float(rotate), w=r_w, h=r_h)

    origin_filename = os.path.basename(file_path)
    out_file = os.path.join(params.out_dir, origin_filename)

    if dst is not None:
        cv2.imwrite(out_file, dst)
        return True
    return False


def crop(img, rect: str, w, h):
    # 裁剪,
    # 如果没有改变，返回None
    if is_empty(rect):
        return None, None, None
    r = rect.split(',')
    if len(r) != 4:
        return None, None, None

    left = int(r[0])
    top = int(r[1])
    right = int(r[2])
    bottom = int(r[3])

    if left < 0:
        left = 0
    if right > w:
        right = w
    if top < 0:
        top = 0
    if bottom > h:
        bottom = h

    if left == 0 and top == 0 and right == w and bottom == h:
        return None, None, None

    return img[top:bottom, left:right], right - left, bottom - top


def rot_degree(img, degree, w, h):
    center = (w / 2, h / 2)

    M = cv2.getRotationMatrix2D(center, degree, 1)
    top_right = np.array((w - 1, 0)) - np.array(center)
    bottom_right = np.array((w - 1, h - 1)) - np.array(center)
    top_right_after_rot = M[0:2, 0:2].dot(top_right)
    bottom_right_after_rot = M[0:2, 0:2].dot(bottom_right)
    new_width = max(int(abs(bottom_right_after_rot[0] * 2) + 0.5), int(abs(top_right_after_rot[0] * 2) + 0.5))
    new_height = max(int(abs(top_right_after_rot[1] * 2) + 0.5), int(abs(bottom_right_after_rot[1] * 2) + 0.5))
    offset_x = (new_width - w) / 2
    offset_y = (new_height - h) / 2
    M[0, 2] += offset_x
    M[1, 2] += offset_y
    dst = cv2.warpAffine(img, M, (new_width, new_height))
    return dst, new_width, new_height