utils.py

import numpy as np
import cv2
from PIL import Image


def rotate_cut_img(im, degree, x_center, y_center, w, h, leftAdjust=False, rightAdjust=False, alph=0.2):
    # degree_ = degree * 180.0 / np.pi
    # print(degree_)
    right = 0
    left = 0
    if rightAdjust:
        right = 1
    if leftAdjust:
        left = 1
    
    box = (max(1, x_center - w / 2 - left * alph * (w / 2))
           , y_center - h / 2,  # ymin
           min(x_center + w / 2 + right * alph * (w / 2), im.size[0] - 1)
           , y_center + h / 2)  # ymax
    
    newW = box[2] - box[0]
    newH = box[3] - box[1]
    tmpImg = im.rotate(degree, center=(x_center, y_center)).crop(box)
    
    return tmpImg, newW, newH


def crop_rect(img, rect, alph=0.15):
    img = np.asarray(img)
    # get the parameter of the small rectangle
    # print("rect!")
    # print(rect)
    center, size, angle = rect[0], rect[1], rect[2]
    min_size = min(size)
    
    if angle > -45:
        center, size = tuple(map(int, center)), tuple(map(int, size))
        # angle-=270
        size = (int(size[0] + min_size * alph), int(size[1] + min_size * alph))
        height, width = img.shape[0], img.shape[1]
        M = cv2.getRotationMatrix2D(center, angle, 1)
        # size = tuple([int(rect[1][1]), int(rect[1][0])])
        img_rot = cv2.warpAffine(img, M, (width, height))
        # cv2.imwrite("debug_im/img_rot.jpg", img_rot)
        img_crop = cv2.getRectSubPix(img_rot, size, center)
    else:
        center = tuple(map(int, center))
        size = tuple([int(rect[1][1]), int(rect[1][0])])
        size = (int(size[0] + min_size * alph), int(size[1] + min_size * alph))
        angle -= 270
        height, width = img.shape[0], img.shape[1]
        M = cv2.getRotationMatrix2D(center, angle, 1)
        img_rot = cv2.warpAffine(img, M, (width, height))
        # cv2.imwrite("debug_im/img_rot.jpg", img_rot)
        img_crop = cv2.getRectSubPix(img_rot, size, center)
    img_crop = Image.fromarray(img_crop)
    return img_crop


def draw_bbox(img_path, result, color=(255, 0, 0), thickness=2):
    if isinstance(img_path, str):
        img_path = cv2.imread(img_path)
        # img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
    img_path = img_path.copy()
    for point in result:
        point = point.astype(int)
        cv2.line(img_path, tuple(point[0]), tuple(point[1]), color, thickness)
        cv2.line(img_path, tuple(point[1]), tuple(point[2]), color, thickness)
        cv2.line(img_path, tuple(point[2]), tuple(point[3]), color, thickness)
        cv2.line(img_path, tuple(point[3]), tuple(point[0]), color, thickness)
    return img_path


def sort_box(boxs):
    res = []
    for box in boxs:
        # box = [x if x>0 else 0 for x in box ]
        x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
        newBox = [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
        # sort x
        newBox = sorted(newBox, key=lambda x: x[0])
        x1, y1 = sorted(newBox[:2], key=lambda x: x[1])[0]
        index = newBox.index([x1, y1])
        newBox.pop(index)
        newBox = sorted(newBox, key=lambda x: -x[1])
        x4, y4 = sorted(newBox[:2], key=lambda x: x[0])[0]
        index = newBox.index([x4, y4])
        newBox.pop(index)
        newBox = sorted(newBox, key=lambda x: -x[0])
        x2, y2 = sorted(newBox[:2], key=lambda x: x[1])[0]
        index = newBox.index([x2, y2])
        newBox.pop(index)
        
        newBox = sorted(newBox, key=lambda x: -x[1])
        x3, y3 = sorted(newBox[:2], key=lambda x: x[0])[0]
        
        res.append([x1, y1, x2, y2, x3, y3, x4, y4])
    return res


def solve(box):
    """
    绕 cx,cy点 w,h 旋转 angle 的坐标
    x = cx-w/2
    y = cy-h/2
    x1-cx = -w/2*cos(angle) +h/2*sin(angle)
    y1 -cy= -w/2*sin(angle) -h/2*cos(angle)

    h(x1-cx) = -wh/2*cos(angle) +hh/2*sin(angle)
    w(y1 -cy)= -ww/2*sin(angle) -hw/2*cos(angle)
    (hh+ww)/2sin(angle) = h(x1-cx)-w(y1 -cy)

    """
    x1, y1, x2, y2, x3, y3, x4, y4 = box[:8]
    cx = (x1 + x3 + x2 + x4) / 4.0
    cy = (y1 + y3 + y4 + y2) / 4.0
    w = (np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) + np.sqrt((x3 - x4) ** 2 + (y3 - y4) ** 2)) / 2
    h = (np.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2) + np.sqrt((x1 - x4) ** 2 + (y1 - y4) ** 2)) / 2
    
    sinA = (h * (x1 - cx) - w * (y1 - cy)) * 1.0 / (h * h + w * w) * 2
    angle = np.arcsin(sinA)
    return angle, w, h, cx, cy


def sorted_boxes(dt_boxes):
    """
    Sort text boxes in order from top to bottom, left to right
    args:
        dt_boxes(array):detected text boxes with shape [4, 2]
    return:
        sorted boxes(array) with shape [4, 2]
    """
    num_boxes = dt_boxes.shape[0]
    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
    _boxes = list(sorted_boxes)

    for i in range(num_boxes - 1):
        if abs(_boxes[i+1][0][1] - _boxes[i][0][1]) < 10 and \
            (_boxes[i + 1][0][0] < _boxes[i][0][0]):
            tmp = _boxes[i]
            _boxes[i] = _boxes[i + 1]
            _boxes[i + 1] = tmp
    return _boxes


def get_rotate_crop_image(img, points):

    left = int(np.min(points[:, 0]))
    right = int(np.max(points[:, 0]))
    top = int(np.min(points[:, 1]))
    bottom = int(np.max(points[:, 1]))
    img_crop = img[top:bottom, left:right, :].copy()
    points[:, 0] = points[:, 0] - left
    points[:, 1] = points[:, 1] - top
    img_crop_width = int(np.linalg.norm(points[0] - points[1]))
    img_crop_height = int(np.linalg.norm(points[0] - points[3]))
    pts_std = np.float32([[0, 0], [img_crop_width, 0],\
        [img_crop_width, img_crop_height], [0, img_crop_height]])

    M = cv2.getPerspectiveTransform(points, pts_std)
    dst_img = cv2.warpPerspective(
        img_crop,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE)
    dst_img_height, dst_img_width = dst_img.shape[0:2]
    if dst_img_height * 1.0 / dst_img_width >= 1.5:
        dst_img = np.rot90(dst_img)
    return dst_img

def get_crop_image(img, points):

    left = int(np.min(points[:, 0]))
    right = int(np.max(points[:, 0]))
    top = int(np.min(points[:, 1]))
    bottom = int(np.max(points[:, 1]))
    img_crop = img[top:bottom, left:right, :].copy()
    
    img_height, img_width = img_crop.shape[0:2]
    if img_height * 1.0 / img_width >= 1.5:
        img_crop = np.rot90(img_crop)
    return img_crop

def get_hsv_range(colorlist):
    hsv_list = []
    if  not isinstance(colorlist,list):
        colorlist = colorlist.split(' · ')

    for color_int in colorlist:
        color_int = int(color_int,16)
        b = (color_int & 255)/255.0
        g = ((color_int >> 8) & 255)/255.0
        r = ((color_int >> 16) & 255)/255.0

        mx = max(r, g, b)
        mn = min(r, g, b)
        m = mx-mn
        if mx == mn:
            h = 0
        elif mx == r:
            if g >= b:
                h = ((g-b)/m)*60
            else:
                h = ((g-b)/m)*60 + 360
        elif mx == g:
            h = ((b-r)/m)*60 + 120
        elif mx == b:
            h = ((r-g)/m)*60 + 240
        if mx == 0:
            s = 0
        else:
            s = m/mx
        v = mx
        
        H = h / 2
        S = s * 255.0
        V = v * 255.0
        hsv_list.append([H,S,V])

    hsv_arr = np.array(hsv_list)
    hsvLower = hsv_arr.min(0)
    hsvUpper = hsv_arr.max(0)
    return hsvLower,hsvUpper