crop_detect_rotate_id_card_yolact.py

import cv2
import math
import numpy as np
import ocr_v2.config.constants as const
from ocr_v2.process.PretrainedModel import PretrainedModel
from ocr_v2.model.yolact.yolact_eval import yolact_predict

models = PretrainedModel()


def angle(vector_1, vector_2):
    unit_vector_1 = vector_1 / np.linalg.norm(vector_1)
    unit_vector_2 = vector_2 / np.linalg.norm(vector_2)
    dot_product = np.dot(unit_vector_1, unit_vector_2)
    angle_rad = np.arccos(dot_product)
    return angle_rad


def distance(vector_1, vector_2):
    distance_res = np.linalg.norm(vector_2 - vector_1)
    return distance_res


def perp(a):
    b = np.zeros_like(a)
    b[0] = -a[1]
    b[1] = a[0]
    return b


def intersect(a1, a2, b1, b2):
    da = a2 - a1
    db = b2 - b1
    dp = a1 - b1
    dap = perp(da)
    denom = np.dot(dap, db)
    num = np.dot(dap, dp)
    return (num / denom.astype(float)) * db + b1


def detect_crop_rotate_card_yolact(img):

    model_crop = models.yolactModel
    model_classify = models.classifyCardSideModel

    corners = []
    all_distances = []
    distances = []
    rect = []
    min_rect_vec = []
    edges = []
    start = 0
    new_rect = [0, 0, 0, 0]
    angle_threshold = 3
    num_edge = 0

    img = img.astype(float)
    _, _, msk, _ = yolact_predict(model_crop, img, 0.9)
    msk = msk[0]
    contours, hierarchy = cv2.findContours(msk.astype('uint8')*255, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    ratio_area = np.sum(msk.astype('uint8')) / (img.shape[0]*img.shape[1])
    max_contour = max(contours, key=cv2.contourArea)
    convex_hull = cv2.convexHull(max_contour)
    min_rect = cv2.minAreaRect(convex_hull)
    min_rect = np.int0(cv2.boxPoints(min_rect))
    min_width = min(distance(min_rect[0], min_rect[1]), distance(min_rect[1], min_rect[2]))
    for i in range(4):
        min_rect_vec.append(min_rect[i] - min_rect[(i+1) % 4])
    for i, point in enumerate(convex_hull):
        if angle(convex_hull[i - 1][0] - point[0],
                 convex_hull[(i + 1) % len(convex_hull)][0] - point[0]) < angle_threshold:
            corners.append(point)
    for i, point in enumerate(corners):
        all_distances.append(distance(point[0], corners[(i+1) % len(corners)][0]))
        min_angle = np.pi / 2
        min_dis_to_rect = img.shape[0]
        current_vec = point[0] - corners[(i + 1) % len(corners)][0]
        for k in range(4):
            ang = angle(current_vec, min_rect_vec[k])
            if ang < min_angle or (ang > np.pi/2 and ang - np.pi/2 < min_angle):
                min_angle = ang if ang < np.pi / 2 else ang - np.pi / 2
            dis1 = distance(point[0], min_rect[k])
            dis2 = distance(corners[(i + 1) % len(corners)][0], min_rect[k])
            min_dis_to_rect = dis1 if min_dis_to_rect > dis1 else min_dis_to_rect
            min_dis_to_rect = dis2 if min_dis_to_rect > dis2 else min_dis_to_rect
        if min_angle < np.pi * (10 / 180) or min_dis_to_rect <= min_width / 12:
            edges.append(i)
            distances.append(distance(point[0], corners[(i + 1) % len(corners)][0]))
    print(distances)
    edges = np.array(edges)
    distances = np.array(distances)
    all_distances = np.array(all_distances)
    max_dis_arg = all_distances.argsort()[-1]
    max_args = distances.argsort()
    for i in range(len(max_args)):
        dis_ratio = distances[max_args[len(max_args) - i - 1]] / all_distances[max_dis_arg]
        if dis_ratio < 1 / 10:
            num_edge = i
            break
    if num_edge < 4:
        edges = np.array(range(len(corners)))
        max_args = all_distances.argsort()
    max_args = np.sort(max_args[-4:])
    for i in range(4):
        if edges[max_args[(i+1) % 4]] == max_dis_arg:
            start = i
        point = intersect(corners[edges[max_args[i]]][0],
                          corners[(edges[max_args[i]] + 1) % len(corners)][0],
                          corners[edges[max_args[(i+1) % 4]]][0],
                          corners[(edges[max_args[(i+1) % 4]] + 1) % len(corners)][0])
        rect.append(point.astype(float))
    for i in range(4):
        new_rect[(4 - start + i) % 4] = rect[i]
    dst = np.array([[0, 0], [1026 - 1, 0], [1026 - 1, 640 - 1], [0, 640 - 1]], dtype="float32")
    m = cv2.getPerspectiveTransform(np.array(new_rect, dtype="float32"), dst)
    warped = cv2.warpPerspective(img, m, (1026, 640))
    warped = warped.astype('uint8')

    dst_copy = warped.copy()
    dst_copy = cv2.resize(dst_copy, (256, 160))
    dst_copy = np.expand_dims(dst_copy, 0)
    y_pred = model_classify.predict(dst_copy)
    y_pred = np.argmax(y_pred, axis=1)
    y_pred = y_pred.astype('int')
    if y_pred == 1 or y_pred == 3:
        warped = warped[::-1, ::-1]
    if y_pred == 0 or y_pred == 1:
        return warped, const.SIDE_FRONT, ratio_area
    if y_pred == 2 or y_pred == 3:
        return warped, const.SIDE_BACK, ratio_area