import cv2
import numpy as np


# https://www.cnblogs.com/meikaiyuan/p/16805539.html

def mid_line_fit(contours_, mask, raw):
    middle_s = []

    for i in range(0, 143):
        count = 0
        sum_ = 0
        for j in range(0, 255):
            if contours_[1][j][0][1] == i:
                sum_ += contours_[1][j][0][0]
                count += 1
        middle = sum_ / count
        print(count)
        middle_s.append([[middle, i]])

    contours_middle = np.array(middle_s).astype(np.int32)

    cv2.drawContours(mask, contours_middle, -1, (0, 255, 0), 2)

    output = cv2.fitLine(contours_middle, cv2.DIST_L2, 0, 0.01, 0.01)
    k = output[1] / output[0]
    b = output[3] - k * output[2]
    x_0 = -b / k
    x_480 = (479 - b) / k
    cv2.line(mask, (int(x_0), 0), (int(x_480), 479), (255, 0, 0), 4)
    cv2.line(raw, (int(x_0), 0), (int(x_480), 479), (255, 0, 0), 4)

    cv2.imshow("img_c", raw)
    cv2.imshow("line", mask)
    # cv2.imshow("mask", mask_c)
    cv2.waitKey(0)


def read_image(path):
    raw = cv2.imread(path)
    # 设定颜色HSV范围，假定为红色
    color_lower = np.array([0, 148, 116])
    color_upper = np.array([50, 220, 203])

    # 将图像转化为HSV格式
    hsv = cv2.cvtColor(raw, cv2.COLOR_BGR2HSV)

    # 去除颜色范围外的其余颜色
    mask = cv2.inRange(hsv, color_lower, color_upper)
    # cv2.imshow("mask", mask)

    # 二值化操作
    ret, binary = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)

    # 膨胀操作，因为是对线条进行提取定位，所以腐蚀可能会造成更大间隔的断点，将线条切断，因此仅做膨胀操作
    kernel = np.ones((5, 5), np.uint8)
    dilation = cv2.dilate(binary, kernel, iterations=1)
    #
    # img2 = cv2.bitwise_and(img, img, mask=mask)
    # cv2.imshow('image', dilation)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # 获取图像轮廓坐标，其中contours为坐标值，此处只检测外形轮廓
    contours, hierarchy = cv2.findContours(dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.approxPolyDP(contours[0], 0.01, True)
    print(contours)
    mid_line_fit(contours, mask, raw)




if __name__ == '__main__':
    read_image("../data/004.png")