提取图片中的红色面积

extract_red.py

import cv2
import numpy as np
import sys
import os


def extract_red(infile, outfile):
    print(infile)
    img = cv2.imread(infile)
    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # lower mask (0-10)
    lower_red = np.array([0, 50, 50])
    upper_red = np.array([10, 255, 255])
    mask0 = cv2.inRange(img_hsv, lower_red, upper_red)

    # upper mask (170-180)
    lower_red = np.array([170, 50, 50])
    upper_red = np.array([180, 255, 255])
    mask1 = cv2.inRange(img_hsv, lower_red, upper_red)

    # join my masks
    mask = mask0 + mask1
    red = np.count_nonzero(mask)
    total = mask.size
    text1 = "红色面积 / 总面积 = 红色比例"
    text2 = "%d / %d = %.2lf" % (red, total, red / total)
    print(text1)
    print(text2)

    # set my output img to zero everywhere except my mask
    output_img = img.copy()
    output_img[np.where(mask == 0)] = 0

    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(output_img, text2, (10, 50), font, 1, (255, 255, 255), 2,
                cv2.LINE_AA)

    # or your HSV image, which I *believe* is what you want
    output_hsv = img_hsv.copy()
    output_hsv[np.where(mask == 0)] = 0

    cv2.imwrite(outfile, output_img)


if __name__ == '__main__':
    extract_red(sys.argv[1], os.path.splitext(sys.argv[1])[0] + '_red.jpg')