detect_color.py 文件源码

def checkRedPixel(inputPath):
    """
    docstring
    """
    im = cv2.imread(inputPath, 1)
    # Convert BGR to HSV
    hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
    # define range of red color in HSV
    lower_red = np.array([0, 100, 100])
    upper_red = np.array([10, 255, 255])
    # Threshold the HSV image to get only red colors
    mask1 = cv2.inRange(hsv, lower_red, upper_red)
    # define range of red color in HSV
    lower_red = np.array([170, 100, 100])
    upper_red = np.array([180, 255, 255])
    # Threshold the HSV image to get only red colors
    mask2 = cv2.inRange(hsv, lower_red, upper_red)
    mask = mask1 + mask2
    # Bitwise-AND mask and original image
    res = cv2.bitwise_and(im, im, mask=mask)
    # save temp image
    cv2.imwrite(os.path.join(TEMP, 'temp.png'), res)
    ct = ColorThief(os.path.join(TEMP, 'temp.png'))
    palette = ct.get_palette(color_count=5)
    for p in palette:
        r = p[0]
        g = p[1]
        b = p[2]
        bgr = np.uint8([[[p[2], p[1], p[0]]]])
        hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
        h = hsv[0][0][0]
        s = hsv[0][0][1]
        v = hsv[0][0][2]
        if ((h >= 0 and h <= 10) and (s >= 100 and s <= 255) and (v >= 100 and v <= 255)) or ((h >= 170 and h <= 180) and (s >= 100 and s <= 255) and (v >= 100 and v <= 255)):
            return True
            break