y2020/vision/galactic_search_config.py

#!/usr/bin/python3

# Creates a UI for a user to select the regions in a camera image where the balls could be placed
# for each field layout.
# After the balls have been placed on the field and they submit the regions,
# galactic_search_path.py 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 galactic_search_path import *

import getopt
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
import sys

_num_rects = 3 # can be 3 or 2, can be specified in commang line arg

setup_if_pi()

_path = Path(Letter.kA, Alliance.kRed, [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(txt):
    txt.set_text("Click on top left point and bottom right point for rect #%u" % (_rect_index + 1))
    txt.set_color(_path.alliance.value)


def on_confirm(event):
    global _rect_index
    if _path.rects[_rect_index].x1 != None and _path.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(_txt)
            _path.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()
        _path.rects[_rect_index].x1 = None
        _path.rects[_rect_index].y1 = None
        _path.rects[_rect_index].x2 = None
        _path.rects[_rect_index].y2 = None
        plt.show()

def on_submit(event):
    plt.close("all")
    dict = None
    with open(RECTS_JSON_PATH, 'r') as rects_json:
        dict = json.load(rects_json)
    if _path.letter.name not in dict:
        dict[_path.letter.name] = {}
    if _path.alliance.name not in dict[_path.letter.name]:
        dict[_path.letter.name][_path.alliance.name] = []
    dict[_path.letter.name][_path.alliance.name] = [rect.to_list() for rect in _path.rects]
    with open(RECTS_JSON_PATH, 'w') as rects_json:
        json.dump(dict, rects_json, indent = 2)

# 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 gets called for each click of the rectangle corners,
    # but also gets called when the user clicks on the Submit button.
    # At that time _rect_index will equal the length of rects, and so we'll ignore that click.
    # If it checked the points of the rect at _rect_index, a list out of bounds exception would be thrown.
    # 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 _path.rects[_rect_index].x1 == None:
            _path.rects[_rect_index].x1, _path.rects[_rect_index].y1 = int(event.xdata), int(event.ydata)
        elif _path.rects[_rect_index].x2 == None:
            _path.rects[_rect_index].x2, _path.rects[_rect_index].y2 = int(event.xdata), int(event.ydata)
            if _path.rects[_rect_index].x2 < _path.rects[_rect_index].x1:
                tmp = _path.rects[_rect_index].x1
                _path.rects[_rect_index].x1 = _path.rects[_rect_index].x2
                _path.rects[_rect_index].x2 = tmp
            if _path.rects[_rect_index].y2 < _path.rects[_rect_index].y1:
                tmp = _path.rects[_rect_index].y1
                _path.rects[_rect_index].y1 = _path.rects[_rect_index].y2
                _path.rects[_rect_index].y2 = tmp

            _img_ax.add_patch(patches.Rectangle((_path.rects[_rect_index].x1, _path.rects[_rect_index].y1),
                _path.rects[_rect_index].x2 - _path.rects[_rect_index].x1,
                _path.rects[_rect_index].y2 - _path.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 setup_ui():
    _img_ax.imshow(capture_img())
    release_stream()

    _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(_txt)
    plt.show()

def main(argv):
    global _num_rects

    glog.setLevel("INFO")
    opts = getopt.getopt(argv[1 : ], "a:l:n:",
                        ["alliance = ", "letter = ", "_num_rects = "])[0]
    for opt, arg in opts:
        if opt in ["-a", "--alliance"]:
            _path.alliance = Alliance.from_value(arg)
        elif opt in ["-l", "--letter"]:
            _path.letter = Letter.from_value(arg.upper())
        elif opt in ["-n", "--_num_rects"] and arg.isdigit():
            _num_rects = int(arg)

    setup_ui()


if __name__ == "__main__":
    main(sys.argv)