pythonProject1/ColorDiscrimination2.0.py at master · Yc-Dlan/pythonProject1 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
from os.path import split

import cv2 as cv
import numpy as np
#import matplotlib.pyplot as plt
#from imutils import contours

def enhanced_brightness_image(image,a,b):
    adjusted_image=cv.convertScaleAbs(image,alpha=a,beta=b)
    return adjusted_image

def preprocess_image(image):
    image_ycrcb=cv.cvtColor(image,cv.COLOR_BGR2YCrCb)
    channels=cv.split(image_ycrcb)
    clahe=cv.createCLAHE(clipLimit=2.0,tileGridSize=(3,3))
    clahe.apply(channels[0],channels[0])
    cv.merge(channels,image_ycrcb)
    cv.cvtColor(image_ycrcb,cv.COLOR_YCR_CB2BGR,image)
    return image

def color_discrimination(frame):        #颜色识别函数
    frame_gsblur=cv.GaussianBlur(frame,(5,5),0)
    frame_hsv=cv.cvtColor(frame_gsblur,cv.COLOR_BGR2HSV)

    red_mask_0=cv.inRange(frame_hsv,lower_red_0,upper_red_0)
    red_mask_1=cv.inRange(frame_hsv,lower_red_1,upper_red_1)
    blue_mask_0=cv.inRange(frame_hsv,lower_blue_0,upper_blue_0)
    yellow_mask_0=cv.inRange(frame_hsv,lower_yellow_0,upper_yellow_0)
    green_mask_0=cv.inRange(frame_hsv,lower_green_0,upper_green_0)
    purple_mask_0=cv.inRange(frame_hsv,lower_purple_0,upper_purple_0)

    kernel = np.ones((3,3), np.uint8)
    red_mask_0 = cv.morphologyEx(red_mask_0,cv.MORPH_OPEN, kernel)
    red_mask_1 = cv.morphologyEx(red_mask_1,cv.MORPH_OPEN, kernel)
    blue_mask_0 = cv.morphologyEx(blue_mask_0,cv.MORPH_OPEN, kernel)
    green_mask_0 = cv.morphologyEx(green_mask_0,cv.MORPH_OPEN, kernel)
    yellow_mask_0 = cv.morphologyEx(yellow_mask_0,cv.MORPH_OPEN,kernel)
    purple_mask_0 = cv.morphologyEx(purple_mask_0,cv.MORPH_OPEN,kernel)

    contours, _=cv.findContours(red_mask_0+red_mask_1+blue_mask_0+green_mask_0+yellow_mask_0+purple_mask_0,cv.RETR_EXTERNAL,cv.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        x,y,w,h=cv.boundingRect(contour)
        color=""
        if cv.contourArea(contour)>2000:
            if np.any(red_mask_0[y:y + h, x:x + w]):
                color = "red"
                print(color)
                print('center_x',x+w/2)
                print('center_y',y+h/2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            elif np.any(red_mask_1[y:y + h, x:x + w]):
                color = "red"
                print(color)
                print('center_x', x + w / 2)
                print('center_y', y + h / 2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            elif np.any(blue_mask_0[y:y + h, x:x + w]):
                color = "blue"
                print(color)
                print('center_x', x + w / 2)
                print('center_y', y + h / 2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
            elif np.any(green_mask_0[y:y + h, x:x + w]):
                color = "green"
                print(color)
                print('center_x', x + w / 2)
                print('center_y', y + h / 2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            elif np.any(yellow_mask_0[y:y + h, x:x + w]):
                color = "yellow"
                print(color)
                print('center_x', x + w / 2)
                print('center_y', y + h / 2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 255), 2)
            elif np.any(purple_mask_0[y:y + h, x:x + w]):
                color = "purple"
                print(color)
                print('center_x', x + w / 2)
                print('center_y', y + h / 2)
                cv.rectangle(frame, (x, y), (x + w, y + h), (128, 0, 128), 2)
            cv.putText(frame, color, (x, y - 10), cv.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)

    return frame

#主函数

lower_red_0 = np.array([0, 100, 100])
upper_red_0 = np.array([6, 255, 255])
lower_red_1 = np.array([160, 100, 100])
upper_red_1 = np.array([170, 255, 255])
lower_blue_0 = np.array([110, 60, 60])
upper_blue_0 = np.array([128, 255, 255])
lower_yellow_0 = np.array([25, 100, 100])
upper_yellow_0 = np.array([35, 255, 255])
lower_green_0 = np.array([45, 150, 150])
upper_green_0 = np.array([85, 255, 255])
lower_purple_0 = np.array([130, 100, 100])
upper_purple_0 = np.array([150, 255, 255])

cap=cv.VideoCapture("image/VID_20240922_172258.mp4")
cap.set(cv.CAP_PROP_FRAME_WIDTH,1280)
cap.set(cv.CAP_PROP_FRAME_HEIGHT,720)

while cap.isOpened():
    ret, frame = cap.read()
    if ret:
        frame=cv.resize(frame,(1280,720))
        frame_brighter=enhanced_brightness_image(frame,2.2,50)
        frame_enhanced=preprocess_image(frame_brighter)
        result = color_discrimination(frame_enhanced)
        cv.imshow("Color Discrimination", result)

        if cv.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
cv.destroyAllWindows()