SelfDrivingCar/ColorDetectionCode.py at main · haladayyeh/SelfDrivingCar · 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
import cv2 as cv
import numpy as np
import imutils

cap = cv.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)
cnt = 0

color_ranges = {
    'gray': ([0, 0, 40], [180, 18, 230]),
    'yellow': ([25, 50, 70], [35, 255, 255]),
    'orange': ([10, 50, 70], [24, 255, 255]),
    'blue': ([90, 50, 70], [128, 255, 255]),
    'green': ([36, 50, 70], [89, 255, 255]),
    'red': ([125, 50, 70], [180, 255, 255])
}

def create_mask(hsv, color_name):
    lower, upper = color_ranges[color_name]
    return cv.inRange(hsv, np.array(lower), np.array(upper))

def process_contours(frame, contours, color, min_area, movement_threshold=None):
    global cnt
    color_bgr = {
        'green': (0, 255, 0),
        'red': (0, 0, 255),
        'blue': (255, 0, 0),
        'yellow': (0, 255, 255),
        'orange': (0, 165, 255),
        'gray': (128, 128, 128)
    }

    for c in contours:
        area = cv.contourArea(c)

        if area > min_area:
            if color == 'green' and area > 15000:
                cnt += 1

            cv.drawContours(frame, [c], -1, color_bgr.get(color, (255, 255, 255)), 2)

            M = cv.moments(c)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                cv.circle(frame, (cx, cy), 7, (255, 255, 255), -1)
                cv.putText(frame, color, (cx-20, cy-20), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

                if movement_threshold and area < 45000:
                    if color == 'green' and cx < 800:
                        print("move to left")
                    elif color == 'red' and cx > 200:
                        print("move to right")

while True:
    ret, frame = cap.read()
    if not ret:
        break

    frame = cv.resize(frame, (1000, 750))
    blur = cv.GaussianBlur(frame, (15, 15), 0)
    hsv = cv.cvtColor(blur, cv.COLOR_BGR2HSV)

    masks = {}
    contours = {}

    for color_name in color_ranges.keys():
        masks[color_name] = create_mask(hsv, color_name)
        cnts = cv.findContours(masks[color_name], cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
        contours[color_name] = imutils.grab_contours(cnts)

    process_contours(frame, contours['green'], 'green', 15000, True)
    process_contours(frame, contours['red'], 'red', 5000, True)
    process_contours(frame, contours['blue'], 'blue', 5000)
    process_contours(frame, contours['yellow'], 'yellow', 5000)
    process_contours(frame, contours['orange'], 'orange', 5000)
    process_contours(frame, contours['gray'], 'gray', 5000)

    cv.imshow("result", frame)

    key = cv.waitKey(5) & 0xFF
    if key == 27 or key == ord('q'):
        break

cap.release()
cv.destroyAllWindows()
print(f"Total green objects detected: {cnt}")