Vision-2023/OptimisedCode.py at main · Decatur-Robotics/Vision-2023 · 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
132
133
134
135
136
137
138
139
140
141
142
import cv2
from RioComms import RioComms
import keyboard
import time
import threading

print("Initializing.")

rioComms = RioComms("10.40.26.2")

cap = cv2.VideoCapture(1)

cap.set(cv2.CAP_PROP_EXPOSURE, -6)

cameraX = 120
cameraY = 70

kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, ksize=(15, 15))

coneHueMask = None
coneSaturationMask = None
cubeHueMask = None
cubeSaturationMask = None
finalConeMask = None
finalCubeMask = None

def cone_hue_mask(lower, upper):
  global coneHueMask

  coneHueMask = (cv2.dilate(cv2.erode(cv2.medianBlur(cv2.inRange(img, (lower, 0, 0), (upper, 255, 255)), 25), kernel), kernel))

def cone_saturation_mask(lower, upper):
  global coneSaturationMask

  coneSaturationMask = (cv2.dilate(cv2.erode(cv2.medianBlur(cv2.inRange(img, (0, lower, 0), (180, upper, 255)), 25), kernel), kernel))

def cone_mask_merge(mask1, mask2):
  global finalConeMask

  finalConeMask = (cv2.bitwise_and(mask1, mask2))

def cube_hue_mask(lower, upper):
  global cubeHueMask

  cubeHueMask = (cv2.dilate(cv2.erode(cv2.medianBlur(cv2.inRange(img, (lower, 0, 0), (upper, 255, 255)), 25), kernel), kernel))

def cube_saturation_mask(lower, upper):
  global cubeSaturationMask

  cubeSaturationMask = (cv2.dilate(cv2.erode(cv2.medianBlur(cv2.inRange(img, (0, lower, 0), (180, upper, 255)), 25), kernel), kernel))

def cube_mask_merge(mask1, mask2):
  global finalCubeMask

  finalCubeMask = (cv2.bitwise_and(mask1, mask2))

def find_cone_position():
  coneCnts = cv2.findContours(finalConeMask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  coneCntsAreas = []

  for i in coneCnts[0]:
    M = cv2.moments(i)

    coneCntsAreas.append(M["m00"])

  if len(coneCntsAreas) > 0:
    M = cv2.moments(coneCnts[0][coneCntsAreas.index(max(coneCntsAreas))])

    cX = int((M["m10"] / M["m00"]))
    cY = int((M["m01"] / M["m00"]))

    rioComms.send("cones", "Cone X", cX - (cameraX / 2))
    rioComms.send("cones", "Cone Y", cY - (cameraY / 2))
    rioComms.send("cones", "Cone Visible", 1)

  else:
    rioComms.send("cones", "Cone X", 0)
    rioComms.send("cones", "Cone Y", 0)
    rioComms.send("cones", "Cone Visible", 0)

def find_cube_position():
  cubeCnts = cv2.findContours(finalCubeMask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  cubeCntsAreas = []

  for i in cubeCnts[0]:
    M = cv2.moments(i)

    cubeCntsAreas.append(M["m00"])

  if len(cubeCntsAreas) > 0:
    M = cv2.moments(cubeCnts[0][cubeCntsAreas.index(max(cubeCntsAreas))])

    cX = int((M["m10"] / M["m00"]))
    cY = int((M["m01"] / M["m00"]))

    rioComms.send("cubes", "Cube X", cX - (cameraX / 2))
    rioComms.send("cubes", "Cube Y", cY - (cameraY / 2))
    rioComms.send("cubes", "Cube Visible", 1)

  else:
    rioComms.send("cubes", "Cube X", 0)
    rioComms.send("cubes", "Cube Y", 0)
    rioComms.send("cubes", "Cube Visible", 0)

print("Running loop.")

while True:
  loopStartTime = time.time()

  if keyboard.is_pressed("z"):
    exit()

  _, img = cap.read()

  img = cv2.cvtColor(cv2.resize(img, (cameraX, cameraY)), cv2.COLOR_BGR2HSV)

  coneHueMaskThread = threading.Thread(target=cone_hue_mask, args=(16, 25))
  coneSaturationMaskThread = threading.Thread(target=cone_saturation_mask, args=(180, 255))

  cubeHueMaskThread = threading.Thread(target=cube_hue_mask, args=(115, 135))
  cubeSaturationMaskThread = threading.Thread(target=cube_saturation_mask, args=(60, 230))

  coneHueMaskThread.start()
  coneSaturationMaskThread.start()
  cubeHueMaskThread.start()
  cubeSaturationMaskThread.start()

  coneHueMaskThread.join()
  coneSaturationMaskThread.join()
  cubeHueMaskThread.join()
  cubeSaturationMaskThread.join()

  coneMaskMergeThread = threading.Thread(target=cone_mask_merge, args=(coneHueMask, coneSaturationMask))
  cubeMaskMergeThread = threading.Thread(target=cube_mask_merge, args=(cubeHueMask, cubeSaturationMask))

  coneMaskMergeThread.start()
  cubeMaskMergeThread.start()

  coneMaskMergeThread.join()
  cubeMaskMergeThread.join()

  print("Looped in", (time.time() - loopStartTime), "seconds.")