Rover2019/object_tracking.py at master · technocratsRover/Rover2019 · 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
120
121
122
123
124
125
126
127
128
129
130
131
import cv2 as cv
import numpy as np

cap = cv.VideoCapture(0)                                          # open webcam for capture
openingKernel = cv.getStructuringElement(cv.MORPH_RECT, (10, 10)) # Make a kernel for closing (noise removal)
openingKernelRed = cv.getStructuringElement(cv.MORPH_RECT, (2, 2))

# define range of blue color in HSV
# Range for green in HSV
# lower = 40 40 40
# upper = 80 255 255
# TODO: Tweek these values to detect all ranges of color properly

# Values calculated from: https://alloyui.com/examples/color-picker/hsv
# Limits for green range
lower_green = np.array([37, 80, 40])
upper_green = np.array([95, 255, 255])

# Limits for blue range
lower_blue = np.array([100, 80, 50])
upper_blue = np.array([125, 255, 255])

# Limits for yellow range
lower_yellow = np.array([20, 80, 40])
upper_yellow = np.array([35, 255, 255])

# Limits for red
lower_red = np.array([0, 80, 40])
upper_red = np.array([5, 255, 255])
while True:
  _, frame = cap.read()                      # Take each frame
  hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV) # Take each frame

  # Threshold the HSV image to get only desired colors
  maskGreen = cv.inRange(hsv, lower_green, upper_green)
  maskBlue = cv.inRange(hsv, lower_blue, upper_blue)
  maskYellow = cv.inRange(hsv, lower_yellow, upper_yellow)
  maskRed = cv.inRange(hsv, lower_red, upper_red)


  # Clear noise from mask by opening it, then apply gaussian blur
  maskGreen = cv.morphologyEx(maskGreen, cv.MORPH_OPEN, openingKernel)
  maskGreen = cv.GaussianBlur(maskGreen, (5, 5), 0)

  maskBlue = cv.morphologyEx(maskBlue, cv.MORPH_OPEN, openingKernel)
  maskBlue = cv.GaussianBlur(maskBlue, (5, 5), 0)

  maskYellow = cv.morphologyEx(maskYellow, cv.MORPH_OPEN, openingKernel)
  maskYellow = cv.GaussianBlur(maskYellow, (5, 5), 0)

  # maskRed = cv.morphologyEx(maskRed, cv.MORPH_OPEN, openingKernelRed)
  # maskRed = cv.GaussianBlur(maskRed, (5, 5), 0)

  # Using threshold incase the image is not binary, can be removed
  _, threshGreen = cv.threshold(maskGreen, 100, 255, cv.THRESH_BINARY)
  _, threshBlue = cv.threshold(maskBlue, 100, 255, cv.THRESH_BINARY)
  _, threshYellow = cv.threshold(maskYellow, 100, 255, cv.THRESH_BINARY)
  _, threshRed = cv.threshold(maskRed, 100, 255, cv.THRESH_BINARY)

  _, contoursGreen, _ = cv.findContours(threshGreen, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
  _, contoursBlue, _ = cv.findContours(threshBlue, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
  _, contoursYellow, _ = cv.findContours(threshYellow, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
  _, contoursRed, _ = cv.findContours(threshRed, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)

  # Bitwise-AND mask and original image
  mask = cv.bitwise_or(maskBlue, maskGreen)
  mask = cv.bitwise_or(mask, maskYellow)
  mask=cv.bitwise_or(mask, maskRed)
  res = cv.bitwise_and(frame, frame, mask=mask)

  for c in contoursGreen:
    # To avoid false-positives, the green object should be sufficiently big
    if cv.contourArea(c) > 700:
      # Calculate center of mass to put the text
      M = cv.moments(c)
      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, "green ball", (cX - 10, cY - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
      cv.drawContours(res, [c], -1, (255, 0, 255), 3)

  for c in contoursBlue:
    # To avoid false-positives, the green object should be sufficiently big
    if cv.contourArea(c) > 1000:
      # Calculate center of mass to put the text
      M = cv.moments(c)
      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, "bottle", (cX - 10, cY - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
      cv.drawContours(res, [c], -1, (255, 0, 255), 3)

  for c in contoursYellow:
    # To avoid false-positives, the green object should be sufficiently big
    if cv.contourArea(c) > 1000:
      # Calculate center of mass to put the text
      M = cv.moments(c)
      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, "yellow disk", (cX - 10, cY - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
      cv.drawContours(res, [c], -1, (255, 0, 255), 3)

  for c in contoursRed:
  # To avoid false-positives, the green object should be sufficiently big
    if cv.contourArea(c) > 1000:
      # Calculate center of mass to put the text
      M = cv.moments(c)
      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, "red box", (cX - 10, cY - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
      cv.drawContours(res, [c], -1, (255, 0, 255), 3)
  # Show the image and the result
  cv.imshow('frame', frame)
  # cv.imshow('res', res)
  # cv.imshow('blue', maskBlue)
  # cv.imshow('green', maskGreen)
  # cv.imshow('yellow', maskYellow)
  cv.imshow('red', maskRed)
  # EXIT if 'esc' key is pressed
  k = cv.waitKey(5) & 0xFF
  if k == 27:
    break

cap.release()         # Release the webcam
cv.destroyAllWindows()# Destroy all the windows