All/Python/PaddleOCR/hmOCR/utility.py

import cv2
import numpy as np
from os import path, listdir

__all__ = [
    "sorted_boxes", "get_rotate_crop_image", "get_min_area_rect_crop",
    "get_image_file_list", "check_and_read"
]


def sorted_boxes(dt_boxes):
    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):
        for j in range(i, -1, -1):
            if abs(_boxes[j + 1][0][1] - _boxes[j][0][1]) < 10 and (_boxes[j + 1][0][0] < _boxes[j][0][0]):
                tmp = _boxes[j]
                _boxes[j] = _boxes[j + 1]  # noqa
                _boxes[j + 1] = tmp
            else:
                break
    return _boxes


def get_rotate_crop_image(img, points):
    """
    img_height, img_width = img.shape[0:2]
    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
    """
    assert len(points) == 4, "shape of points must be 4*2"
    img_crop_width = int(max(np.linalg.norm(points[0] - points[1]), np.linalg.norm(points[2] - points[3])))
    img_crop_height = int(max(np.linalg.norm(points[0] - points[3]), np.linalg.norm(points[1] - points[2])))
    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)  # noqa
    dst_img = cv2.warpPerspective(  # noqa
        img, M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC  # noqa
    )
    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_min_area_rect_crop(img, points):
    bounding_box = cv2.minAreaRect(np.array(points).astype(np.int32))  # noqa
    points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])  # noqa

    if points[1][1] > points[0][1]:
        index_a = 0
        index_d = 1
    else:
        index_a = 1
        index_d = 0
    if points[3][1] > points[2][1]:
        index_b = 2
        index_c = 3
    else:
        index_b = 3
        index_c = 2

    box = [points[index_a], points[index_b], points[index_c], points[index_d]]
    crop_img = get_rotate_crop_image(img, np.array(box))
    return crop_img


def _check_image_file(file_path):
    img_end = ("jpg", "bmp", "png", "jpeg", "rgb", "tif", "tiff", "gif")
    return any([file_path.lower().endswith(e) for e in img_end])


def get_image_file_list(img_file):
    images = []
    if img_file is None or not path.exists(img_file):
        raise Exception(f"not found any img file in {img_file}")

    if path.isfile(img_file) and _check_image_file(img_file):
        images.append(img_file)
    elif path.isdir(img_file):
        for single_file in listdir(img_file):
            file_path = path.join(img_file, single_file)
            if path.isfile(file_path) and _check_image_file(file_path):
                images.append(file_path)
    if len(images) == 0:
        raise Exception(f"not found any img file in {img_file}")
    images = sorted(images)
    return images


def check_and_read(img_path):
    if path.basename(img_path)[-3:].lower() == "gif":
        gif = cv2.VideoCapture(img_path)  # noqa
        ret, frame = gif.read()
        if not ret:
            print(f"Cannot read {img_path}. This gif image maybe corrupted.")
            return None, False
        if len(frame.shape) == 2 or frame.shape[-1] == 1:
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)  # noqa
        img = frame[:, :, ::-1]
        return img, True
    return None, False