import numpy as np
import cv2 as cv


def nothing(x):
    pass


# 预测框坐标初始化
targetPos_x = 0
targetPos_y = 0
lastPos_x = 0
lastPos_y = 0

# 创建一个黑色的图像，一个窗口
cv.namedWindow("image", 0)
cv.resizeWindow("image", 1000, 750)
cv.namedWindow('image')

# 创建颜色变化的轨迹栏
cv.createTrackbar('H low', 'image', 14, 179, nothing)
cv.createTrackbar('H high', 'image', 27, 179, nothing)
cv.createTrackbar('S low', 'image', 107, 255, nothing)
cv.createTrackbar('S high', 'image', 199, 255, nothing)
cv.createTrackbar('V low', 'image', 131, 255, nothing)
cv.createTrackbar('V high', 'image', 226, 255, nothing)
cv.createTrackbar('Threshold', 'image', 100000, 700000, nothing)  # 阈值设定
while (1):
    # 读入图片
    img = cv.imread('DSC04716.JPG')
    # 转换颜色空间 BGR 到 HSV
    hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    # 按下键盘退出窗口
    k = cv.waitKey(1) & 0xFF
    if k == 27:
        break
    # 得到轨迹的当前位置
    H_low = cv.getTrackbarPos('H low', 'image')
    H_high = cv.getTrackbarPos('H high', 'image')
    S_low = cv.getTrackbarPos('S low', 'image')
    S_high = cv.getTrackbarPos('S high', 'image')
    V_low = cv.getTrackbarPos('V low', 'image')
    V_high = cv.getTrackbarPos('V high', 'image')
    threshold = cv.getTrackbarPos('Threshold', 'image')

    # HSV 值范围设定
    lower_color = np.array([H_low, S_low, V_low])
    high_color = np.array([H_high, S_high, V_high])
    # 设置HSV的阈值使得只取某种颜色
    mask = cv.inRange(hsv, lower_color, high_color)  # 获取遮罩

    # 将掩膜和图像逐像素相加
    res = cv.bitwise_and(img, img, mask=mask)
    # 查找轮廓
    # CV_RETR_EXTERNAL 只检测最外围轮廓
    # CV_CHAIN_APPROX_NONE 保存物体边界上所有连续的轮廓点到contours向量内
    _, contours, hierarchy = cv.findContours(mask, cv.RETR_EXTERNAL,
                                             cv.CHAIN_APPROX_NONE)
    # 绘制轮廓
    cv.drawContours(res, contours, -1, (0, 0, 255), 20)
    # 排除干扰，计算位置，绘制方框
    x, y, w, h = 0, 0, 0, 0
    for cnt in contours:
        area = cv.contourArea(cnt)
        # print(area)
        if area > threshold:  # 过滤阈值
            x, y, w, h = cv.boundingRect(cnt)
            lastPos_x = targetPos_x
            lastPos_y = targetPos_y
            targetPos_x = int(x + w / 2)
            targetPos_y = int(y + h / 2)
            cv.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 20)
            # cv.circle(img, (targetPos_x, targetPos_y), 20, (0, 0, 255), 20) # 绘制方框中心

    cv.putText(img, "({:0<2d}, {:0<2d})".format(targetPos_x, targetPos_y),
               (150, 150), cv.FONT_HERSHEY_PLAIN, 10, (0, 255, 0), 5)  # 文字
    imgs = np.hstack([img, res])  # 合并两张图片
    cv.imshow('image', imgs)

cv.destroyAllWindows()
cv.imwrite('result.jpg', imgs)