Add galactic search challenge ball detection
In the galactic search challenge for the first skills challenge,
the robot has to determine which path is present based on the
positions of the balls in the field, and run that respective spline.
This code has a UI for selecting the regions to look for balls in,
and uses those regions to determine the path. It sends the path
using aos_send to the /camera channel.
Change-Id: Ifd7aacc75eb830b762e14487bad0720dcc2a3944
diff --git a/y2020/vision/galactic_search_path.py b/y2020/vision/galactic_search_path.py
new file mode 100644
index 0000000..e7f8402
--- /dev/null
+++ b/y2020/vision/galactic_search_path.py
@@ -0,0 +1,159 @@
+#!/usr/bin/python3
+
+from rect import Rect
+import ball_detection
+
+# Creates a UI for a user to select the regions in a camera image where the balls could be placed.
+# After the balls have been placed on the field and they submit the regions,
+# it will take another picture and based on the yellow regions in that picture it will determine where the
+# balls are. This tells us which path the current field is. It then sends the Alliance and Letter of the path
+# with aos_send to the /camera channel for the robot to excecute the spline for that path.
+
+from enum import Enum
+import glog
+import json
+import matplotlib.patches as patches
+import matplotlib.pyplot as plt
+from matplotlib.widgets import Button
+import numpy as np
+import os
+
+class Alliance(Enum):
+ kRed = "red"
+ kBlue = "blue"
+ kUnknown = None
+
+class Letter(Enum):
+ kA = "A"
+ kB = "B"
+
+
+NUM_RECTS = 4
+AOS_SEND_PATH = "bazel-bin/aos/aos_send"
+
+# The minimum percentage of yellow for a region of a image to
+# be considered to have a ball
+BALL_PCT_THRESHOLD = 50
+
+rects = [Rect(None, None, None, None)]
+
+# current index in rects list
+rect_index = 0
+
+fig, img_ax = plt.subplots()
+
+txt = img_ax.text(0, 0, "", size = 10, backgroundcolor = "white")
+
+confirm = Button(plt.axes([0.7, 0.05, 0.1, 0.075]), "Confirm")
+cancel = Button(plt.axes([0.81, 0.05, 0.1, 0.075]), "Cancel")
+submit = Button(plt.axes([0.4, 0.4, 0.1, 0.1]), "Submit")
+
+def draw_txt():
+ alliance = (Alliance.kRed if rect_index % 2 == 0 else Alliance.kBlue)
+ letter = (Letter.kA if rect_index < (NUM_RECTS / 2) else Letter.kB)
+ txt.set_text("Click on top left point and bottom right point for " +
+ alliance.value + ", path " + letter.value)
+ txt.set_color(alliance.value)
+
+
+def on_confirm(event):
+ global rect_index
+ if rects[rect_index].x1 != None and rects[rect_index].x2 != None:
+ confirm.ax.set_visible(False)
+ cancel.ax.set_visible(False)
+ rect_index += 1
+ clear_rect()
+ if rect_index == NUM_RECTS:
+ submit.ax.set_visible(True)
+ else:
+ draw_txt()
+ rects.append(Rect(None, None, None, None))
+ plt.show()
+
+def on_cancel(event):
+ global rect_index
+ if rect_index < NUM_RECTS:
+ confirm.ax.set_visible(False)
+ cancel.ax.set_visible(False)
+ clear_rect()
+ rects[rect_index].x1 = None
+ rects[rect_index].y1 = None
+ rects[rect_index].x2 = None
+ rects[rect_index].y2 = None
+ plt.show()
+
+def on_submit(event):
+ plt.close("all")
+ pcts = ball_detection.pct_yellow(ball_detection.capture_img(), rects)
+ if len(pcts) == len(rects):
+ paths = []
+ for i in range(len(pcts)):
+ alliance = (Alliance.kRed if i % 2 == 0 else Alliance.kBlue)
+ letter = (Letter.kA if i < NUM_RECTS / 2 else Letter.kB)
+ paths.append({"alliance" : alliance.name, "letter" : letter.name})
+ max_index = np.argmax(pcts)
+ path = paths[max_index]
+ # Make sure that exactly one percentage is >= the threshold
+ rects_with_balls = np.where(pcts >= BALL_PCT_THRESHOLD)[0].size
+ glog.info("rects_with_balls: %s" % rects_with_balls)
+ if rects_with_balls != 1:
+ path["alliance"] = Alliance.kUnknown.name
+ glog.warn("More than one ball found, path is unknown" if rects_with_balls > 1 else
+ "No balls found")
+ glog.info("Path is %s" % path)
+ os.system(AOS_SEND_PATH + " --config bazel-bin/y2020/config.json " +
+ "/pi1/camera y2020.vision.GalacticSearchPath '" + json.dumps(path) + "'")
+
+ for j in range(len(pcts)):
+ glog.info("%s: %s%% yellow" % (rects[j], pcts[j]))
+ else:
+ glog.error("Error: len of pcts (%u) != len of rects: (%u)" % (len(pcts), len(rects)))
+
+# Clears rect on screen
+def clear_rect():
+ if len(img_ax.patches) == 0:
+ glog.error("There were no patches found in img_ax")
+ else:
+ img_ax.patches[-1].remove()
+
+def on_click(event):
+ # This will get called when user clicks on Submit button, don't want to override the points on
+ # the last rect. Additionally, the event xdata or ydata will be None if the user clicks out of
+ # the bounds of the axis
+ if rect_index < NUM_RECTS and event.xdata != None and event.ydata != None:
+ if rects[rect_index].x1 == None:
+ rects[rect_index].x1, rects[rect_index].y1 = int(event.xdata), int(event.ydata)
+ elif rects[rect_index].x2 == None:
+ rects[rect_index].x2, rects[rect_index].y2 = int(event.xdata), int(event.ydata)
+ if rects[rect_index].x2 < rects[rect_index].x1:
+ rects[rect_index].x2 = rects[rect_index].x1 + (rects[rect_index].x1 - rects[rect_index].x2)
+ if rects[rect_index].y2 < rects[rect_index].y1:
+ rects[rect_index].y2 = rects[rect_index].y1 + (rects[rect_index].y1 - rects[rect_index].y2)
+
+ img_ax.add_patch(patches.Rectangle((rects[rect_index].x1, rects[rect_index].y1),
+ rects[rect_index].x2 - rects[rect_index].x1, rects[rect_index].y2 - rects[rect_index].y1,
+ edgecolor = 'r', linewidth = 1, facecolor="none"))
+ confirm.ax.set_visible(True)
+ cancel.ax.set_visible(True)
+ plt.show()
+ else:
+ glog.info("Either submitted or user pressed out of the bounds of the axis")
+
+def setup_button(button, on_clicked):
+ button.on_clicked(on_clicked)
+ button.ax.set_visible(False)
+
+def main():
+ glog.setLevel("INFO")
+
+ img_ax.imshow(ball_detection.capture_img())
+
+ fig.canvas.mpl_connect("button_press_event", on_click)
+ setup_button(confirm, on_confirm)
+ setup_button(cancel, on_cancel)
+ setup_button(submit, on_submit)
+ draw_txt()
+ plt.show()
+
+if __name__ == "__main__":
+ main()