判断一个mask是否近似直线

0、前言

通过分割得到一个物体的mask，如何判断其是否近似直线？本文提供了一个解决方案。

1、方法

主要利用mask提取骨架线，对骨架线拟合一条直线，然后利用一些统计量，如真实骨架线和拟合直线的MSE及欧式距离、以及mask本身最小外接矩形的宽高比等。

具体代码如下：

import numpy as np


class Line(object):
    def __init__(self, k, b):
        self.k = k
        self.b = b

    def func(self, x):
        return self.k * x + self.b


def parse_minAreaRect(rect):
    """
    旋转矩形框长边与x轴的夹角，在第一象限时，角度为-90°~0°，在第二象限时，角度为0°~90°；
    """
    w, h = rect[1]
    if w > h:
        angle = int(rect[2])
    else:
        angle = -(90 - int(rect[2]))
    return w, h, angle


class LineClassifier(object):
    def __init__(self, auto_Thres=False, Thres_MSE=4, Thres_ED=50, Thres_short=0.001):
        """
        auto_Thres： 是否使用自适应阈值，如不使用，需给出各个阈值
        Thres_MSE: 均方误差（Mean Squared Error）阈值
        Thres_ED:  欧氏距离（Euclidean Distance）阈值
        Thres_short: 骨骼线最小外接矩形框的短边占原图短边比例的阈值
        """
        self.auto_Thres = auto_Thres
        if not self.auto_Thres:
            # 用户指定阈值
            self.Thres_MSE = Thres_MSE
            self.Thres_ED = Thres_ED
            self.Thres_short = Thres_short

    def adjusted_thres(self, hw, HW):
        # 根据不同的大图尺寸，设定不同的基础阈值，该基础阈值会进行自适应调整
        if min(HW) >2000:
            self.Thres_MSE = 5
            self.Thres_ED = 50
        elif 1080 < min(HW) <= 2000:
            self.Thres_MSE = 2
            self.Thres_ED = 30
        elif min(HW) <= 1080:
            self.Thres_MSE = 1
            self.Thres_ED = 10
        self.Thres_short = 0.008  # 此阈值固定不变
        r = max(hw) / (min(hw) + 1e-5)
        multiple = self.get_multiple(r)
        # 根据此倍数调整阈值
        self.Thres_MSE *= multiple
        self.Thres_ED *= multiple
        # self.Thres_short *= multiple

    @staticmethod
    def get_multiple(r):
        if r < 5:
            multiple = 0.1 * r
        elif 5 <= r < 10:
            multiple = 0.2 * r - 0.5
        elif 10 <= r < 20:
            multiple = 0.5 * r - 3.5
        else:
            multiple = r - 13.5
        return multiple

    def is_straight(self, skeleton_pts, minAreaRect, HW, box_hw):
        """
        skeleton_pts: 裂缝骨骼线坐标点
        s_minAreaRect: 裂缝骨骼线的最小外接矩形框 ——（（x,y）,(w,h), angle）形式
        HW: 原图高宽
        return：是否属于直线
        """
        H, W = HW
        bh, bw = box_hw
        sw, sh, a = parse_minAreaRect(minAreaRect)

        # 利用旋转矩形框的角度与中心点，根据斜截式得到直线
        # 求斜率
        a = 180 - a if a > 0 else abs(a)    # 角度
        k = np.tan(np.pi * a / 180)         # 斜率
        # 根据中心点，求截距
        (cx, cy) = minAreaRect[0]
        cy = bh - cy
        b = cy - k * cx
        # 根据斜截式，得到直线
        line = Line(k, b)

        pts_y, pts_x = skeleton_pts[:, 0], skeleton_pts[:, 1]   # skeleton_pts: 行、列
        pts_y = np.array([bh-p for p in pts_y])                 # 相对于instance，而非原图
        pred_y = np.array(list(map(line.func, pts_x)))

        MSE = round(((pred_y - pts_y) ** 2).mean(), 2)
        ED = round(np.linalg.norm(pred_y - pts_y), 2)
        if self.auto_Thres:
            # 根据裂缝骨骼线的最小外接矩形的宽高信息，自适应调整阈值
            self.adjusted_thres((sh, sw), (H, W))

        cls = (MSE < self.Thres_MSE and ED < self.Thres_ED) or (min(sw, sh) / min(H, W) < self.Thres_short)
        # print("self.Thres_MSE:{}, self.Thres_ED:{}, MSE:{}, ED:{}, size_ratio:{}, cls:{}".format(
        #     self.Thres_MSE, self.Thres_ED, MSE, ED, (min(sw, sh) / min(H, W)), cls))
        return cls