import cv2
import numpy as np

import ch340

flag = False
vidId = 1
cap = cv2.VideoCapture(vidId, cv2.CAP_DSHOW)
while not flag:
    suc, im = cap.read()
    cv2.imshow("image", im)
    k = cv2.waitKey(0)
    if k & 0xFF == ord("q"):
        # cap.release()
        # cv2.destroyAllWindows()
        flag = True
    else:
        vidId += 1
        cap.release()
        cap.open(vidId, cv2.CAP_DSHOW)
        print(f"使用摄像头索引: {vidId}")


base = 30
pump = ch340.CH340()
while True:
    suc, im = cap.read()
    hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
    s = hsv[:, :, 1]
    mask = s > base
    cv2.imshow("mask", im * mask[:, :, np.newaxis])
    tot = mask.shape[0] * mask.shape[1]
    val = np.sum(mask)
    rate = val / tot
    thr = (0.05, 0.2, 0.5)
    ret = ""
    if rate < thr[0]:
        ret = "transport"
    elif rate < thr[1]:
        ret = "middle"
    elif rate < thr[2]:
        ret = "about"
    else:
        ret = "colored"
    im = cv2.putText(
        im,
        f"Rate: {rate:.2},base={base},{ret}",
        (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX,
        1,
        (0, 255, 0),
        2,
    )
    cv2.imshow("image", im)
    k = cv2.waitKey(1)
    if k & 0xFF == ord("w"):
        base += 1
    elif k & 0xFF == ord("s"):
        base -= 1
    elif k & 0xFF == ord("a"):
        pump.push_async(speed=0.1, vol=0.1)
    elif k & 0xFF == ord("q"):
        print("Camera index = ", vidId)
        print("k =", base)
        cap.release()
        cv2.destroyAllWindows()
        break
