| #!/usr/bin/python3 |
| |
| from __future__ import print_function |
| # matplotlib overrides fontconfig locations, so it needs to be imported before gtk. |
| import matplotlib.pyplot as plt |
| import os |
| from frc971.control_loops.python import basic_window |
| from frc971.control_loops.python.color import Color, palette |
| import random |
| import gi |
| import numpy as np |
| |
| gi.require_version('Gtk', '3.0') |
| from gi.repository import Gdk, Gtk |
| import cairo |
| from y2023.control_loops.python.graph_tools import to_theta, to_xy, alpha_blend, shift_angles, get_xy |
| from y2023.control_loops.python.graph_tools import l1, l2, joint_center |
| from y2023.control_loops.python.graph_tools import DRIVER_CAM_POINTS |
| from y2023.control_loops.python import graph_paths |
| |
| from frc971.control_loops.python.basic_window import quit_main_loop, set_color, OverrideMatrix, identity |
| |
| import shapely |
| from shapely.geometry import Polygon |
| |
| |
| def px(cr): |
| return OverrideMatrix(cr, identity) |
| |
| |
| # Draw lines to cr + stroke. |
| def draw_lines(cr, lines): |
| cr.move_to(lines[0][0], lines[0][1]) |
| for pt in lines[1:]: |
| cr.line_to(pt[0], pt[1]) |
| with px(cr): |
| cr.stroke() |
| |
| |
| def draw_px_cross(cr, length_px): |
| """Draws a cross with fixed dimensions in pixel space.""" |
| with px(cr): |
| x, y = cr.get_current_point() |
| cr.move_to(x, y - length_px) |
| cr.line_to(x, y + length_px) |
| cr.stroke() |
| |
| cr.move_to(x - length_px, y) |
| cr.line_to(x + length_px, y) |
| cr.stroke() |
| |
| |
| def angle_dist_sqr(a1, a2): |
| """Distance between two points in angle space.""" |
| return (a1[0] - a2[0])**2 + (a1[1] - a2[1])**2 |
| |
| |
| # Find the highest y position that intersects the vertical line defined by x. |
| def inter_y(x): |
| return np.sqrt((l2 + l1)**2 - (x - joint_center[0])**2) + joint_center[1] |
| |
| |
| # Define min and max l1 angles based on vertical constraints. |
| def get_angle(boundary): |
| h = np.sqrt((l1)**2 - (boundary - joint_center[0])**2) + joint_center[1] |
| return np.arctan2(h, boundary - joint_center[0]) |
| |
| |
| # Rotate a rasterized loop such that it aligns to when the parameters loop |
| def rotate_to_jump_point(points): |
| last_pt = points[0] |
| for pt_i in range(1, len(points)): |
| pt = points[pt_i] |
| delta = last_pt[1] - pt[1] |
| if abs(delta) > np.pi: |
| return points[pt_i:] + points[:pt_i] |
| last_pt = pt |
| return points |
| |
| |
| # shift points vertically by dy. |
| def y_shift(points, dy): |
| return [(x, y + dy) for x, y in points] |
| |
| |
| # Get the closest point to a line from a test pt. |
| def get_closest(prev, cur, pt): |
| dx_ang = (cur[0] - prev[0]) |
| dy_ang = (cur[1] - prev[1]) |
| |
| d = np.sqrt(dx_ang**2 + dy_ang**2) |
| if (d < 0.000001): |
| return prev, np.sqrt((prev[0] - pt[0])**2 + (prev[1] - pt[1])**2) |
| |
| pdx = -dy_ang / d |
| pdy = dx_ang / d |
| |
| dpx = pt[0] - prev[0] |
| dpy = pt[1] - prev[1] |
| |
| alpha = (dx_ang * dpx + dy_ang * dpy) / d / d |
| |
| if (alpha < 0): |
| return prev, np.sqrt((prev[0] - pt[0])**2 + (prev[1] - pt[1])**2) |
| elif (alpha > 1): |
| return cur, np.sqrt((cur[0] - pt[0])**2 + (cur[1] - pt[1])**2) |
| else: |
| return (alpha_blend(prev[0], cur[0], alpha), alpha_blend(prev[1], cur[1], alpha)), \ |
| abs(dpx * pdx + dpy * pdy) |
| |
| |
| def closest_segment(lines, pt): |
| c_pt, c_pt_dist = get_closest(lines[-1], lines[0], pt) |
| for i in range(1, len(lines)): |
| prev = lines[i - 1] |
| cur = lines[i] |
| c_pt_new, c_pt_new_dist = get_closest(prev, cur, pt) |
| if c_pt_new_dist < c_pt_dist: |
| c_pt = c_pt_new |
| c_pt_dist = c_pt_new_dist |
| return c_pt, c_pt_dist |
| |
| |
| # Defining outline of robot |
| class RobotOutline(): |
| |
| def __init__(self, drivetrain_pos, drivetrain_width, joint_center_radius, |
| joint_tower_pos, joint_tower_width, joint_tower_height, |
| driver_cam_pos, driver_cam_width, driver_cam_height): |
| self.drivetrain_pos = drivetrain_pos |
| self.drivetrain_width = drivetrain_width |
| |
| self.joint_center_radius = joint_center_radius |
| self.joint_tower_pos = joint_tower_pos |
| self.joint_tower_width = joint_tower_width |
| self.joint_tower_height = joint_tower_height |
| |
| self.driver_cam_pos = driver_cam_pos |
| self.driver_cam_width = driver_cam_width |
| self.driver_cam_height = driver_cam_height |
| |
| def draw(self, cr): |
| set_color(cr, palette["BLUE"]) |
| |
| cr.move_to(self.drivetrain_pos[0], self.drivetrain_pos[1]) |
| cr.line_to(self.drivetrain_pos[0] + self.drivetrain_width, |
| self.drivetrain_pos[1]) |
| |
| with px(cr): |
| cr.stroke() |
| |
| # Draw joint center |
| cr.arc(joint_center[0], joint_center[1], self.joint_center_radius, 0, |
| 2.0 * np.pi) |
| with px(cr): |
| cr.stroke() |
| |
| # Draw joint tower |
| cr.rectangle(self.joint_tower_pos[0], self.joint_tower_pos[1], |
| self.joint_tower_width, self.joint_tower_height) |
| |
| with px(cr): |
| cr.stroke() |
| |
| # Draw driver cam |
| cr.set_source_rgba(1, 0, 0, 0.5) |
| cr.rectangle(self.driver_cam_pos[0], self.driver_cam_pos[1], |
| self.driver_cam_width, self.driver_cam_height) |
| |
| with px(cr): |
| cr.fill() |
| |
| def draw_theta(self, cr): |
| # TOOD(Max): add theta mode drawing |
| pass |
| |
| |
| DRIVETRAIN_X = -0.490 |
| DRIVETRAIN_Y = 0.184 |
| DRIVETRAIN_WIDTH = 0.980 |
| |
| JOINT_CENTER_RADIUS = 0.173 / 2 |
| |
| JOINT_TOWER_X = -0.252 |
| JOINT_TOWER_Y = DRIVETRAIN_Y |
| JOINT_TOWER_WIDTH = 0.098 |
| JOINT_TOWER_HEIGHT = 0.864 |
| |
| DRIVER_CAM_X = DRIVER_CAM_POINTS[0][0] |
| DRIVER_CAM_Y = DRIVER_CAM_POINTS[0][1] |
| DRIVER_CAM_WIDTH = DRIVER_CAM_POINTS[-1][0] - DRIVER_CAM_POINTS[0][0] |
| DRIVER_CAM_HEIGHT = DRIVER_CAM_POINTS[-1][1] - DRIVER_CAM_POINTS[0][1] |
| |
| |
| class SegmentSelector(basic_window.BaseWindow): |
| |
| def __init__(self, segments): |
| super(SegmentSelector, self).__init__() |
| |
| self.window = Gtk.Window() |
| self.window.set_title("Segment Selector") |
| |
| self.segments = segments |
| |
| self.segment_store = Gtk.ListStore(int, str) |
| |
| for i, segment in enumerate(segments): |
| self.segment_store.append([i, segment.name]) |
| |
| self.segment_box = Gtk.ComboBox.new_with_model_and_entry( |
| self.segment_store) |
| self.segment_box.connect("changed", self.on_combo_changed) |
| self.segment_box.set_entry_text_column(1) |
| |
| self.current_path_index = None |
| |
| self.window.add(self.segment_box) |
| self.window.show_all() |
| |
| def on_combo_changed(self, combo): |
| iter = combo.get_active_iter() |
| |
| if iter is not None: |
| model = combo.get_model() |
| id, name = model[iter][:2] |
| print("Selected: ID=%d, name=%s" % (id, name)) |
| self.current_path_index = id |
| |
| |
| # Create a GTK+ widget on which we will draw using Cairo |
| class ArmUi(basic_window.BaseWindow): |
| |
| def __init__(self, segments): |
| super(ArmUi, self).__init__() |
| |
| self.window = Gtk.Window() |
| self.window.set_title("DrawingArea") |
| |
| self.window.set_events(Gdk.EventMask.BUTTON_PRESS_MASK |
| | Gdk.EventMask.BUTTON_RELEASE_MASK |
| | Gdk.EventMask.POINTER_MOTION_MASK |
| | Gdk.EventMask.SCROLL_MASK |
| | Gdk.EventMask.KEY_PRESS_MASK) |
| self.method_connect("key-press-event", self.do_key_press) |
| self.method_connect("motion-notify-event", self.do_motion) |
| self.method_connect("button-press-event", |
| self._do_button_press_internal) |
| self.method_connect("configure-event", self._do_configure) |
| self.window.add(self) |
| self.window.show_all() |
| |
| self.theta_version = False |
| self.reinit_extents() |
| |
| self.last_pos = to_xy(*graph_paths.points['Neutral'][:2]) |
| self.circular_index_select = 1 |
| |
| # Extra stuff for drawing lines. |
| self.segments = segments |
| self.prev_segment_pt = None |
| self.now_segment_pt = None |
| self.spline_edit = 0 |
| self.edit_control1 = True |
| |
| self.joint_thetas = None |
| self.joint_points = None |
| self.fig = plt.figure() |
| self.axes = [ |
| self.fig.add_subplot(3, 1, 1), |
| self.fig.add_subplot(3, 1, 2), |
| self.fig.add_subplot(3, 1, 3) |
| ] |
| self.fig.subplots_adjust(hspace=1.0) |
| plt.show(block=False) |
| |
| self.index = 0 |
| |
| self.outline = RobotOutline([DRIVETRAIN_X, DRIVETRAIN_Y], |
| DRIVETRAIN_WIDTH, JOINT_CENTER_RADIUS, |
| [JOINT_TOWER_X, JOINT_TOWER_Y], |
| JOINT_TOWER_WIDTH, JOINT_TOWER_HEIGHT, |
| [DRIVER_CAM_X, DRIVER_CAM_Y], |
| DRIVER_CAM_WIDTH, DRIVER_CAM_HEIGHT) |
| |
| self.segment_selector = SegmentSelector(self.segments) |
| self.segment_selector.show() |
| |
| def _do_button_press_internal(self, event): |
| o_x = event.x |
| o_y = event.y |
| x = event.x - self.window_shape[0] / 2 |
| y = self.window_shape[1] / 2 - event.y |
| scale = self.get_current_scale() |
| event.x = x / scale + self.center[0] |
| event.y = y / scale + self.center[1] |
| self.do_button_press(event) |
| event.x = o_x |
| event.y = o_y |
| |
| def _do_configure(self, event): |
| self.window_shape = (event.width, event.height) |
| |
| def redraw(self): |
| if not self.needs_redraw: |
| self.needs_redraw = True |
| self.window.queue_draw() |
| |
| def method_connect(self, event, cb): |
| |
| def handler(obj, *args): |
| cb(*args) |
| |
| self.window.connect(event, handler) |
| |
| def reinit_extents(self): |
| if self.theta_version: |
| self.extents_x_min = -np.pi * 2 |
| self.extents_x_max = np.pi * 2 |
| self.extents_y_min = -np.pi * 2 |
| self.extents_y_max = np.pi * 2 |
| else: |
| self.extents_x_min = -40.0 * 0.0254 |
| self.extents_x_max = 40.0 * 0.0254 |
| self.extents_y_min = -4.0 * 0.0254 |
| self.extents_y_max = 110.0 * 0.0254 |
| |
| self.init_extents( |
| (0.5 * (self.extents_x_min + self.extents_x_max), 0.5 * |
| (self.extents_y_max + self.extents_y_min)), |
| (1.0 * (self.extents_x_max - self.extents_x_min), 1.0 * |
| (self.extents_y_max - self.extents_y_min))) |
| |
| # Handle the expose-event by drawing |
| def handle_draw(self, cr): |
| # use "with px(cr): blah;" to transform to pixel coordinates. |
| if self.segment_selector.current_path_index is not None: |
| self.index = self.segment_selector.current_path_index |
| |
| # Fill the background color of the window with grey |
| set_color(cr, palette["GREY"]) |
| cr.paint() |
| |
| # Draw a extents rectangle |
| set_color(cr, palette["WHITE"]) |
| cr.rectangle(self.extents_x_min, self.extents_y_min, |
| (self.extents_x_max - self.extents_x_min), |
| self.extents_y_max - self.extents_y_min) |
| cr.fill() |
| |
| if self.theta_version: |
| # Draw a filled white rectangle. |
| set_color(cr, palette["WHITE"]) |
| cr.rectangle(-np.pi, -np.pi, np.pi * 2.0, np.pi * 2.0) |
| cr.fill() |
| |
| set_color(cr, palette["BLUE"]) |
| for i in range(-6, 6): |
| cr.move_to(-40, -40 + i * np.pi) |
| cr.line_to(40, 40 + i * np.pi) |
| with px(cr): |
| cr.stroke() |
| |
| set_color(cr, Color(0.0, 1.0, 0.2)) |
| cr.move_to(self.last_pos[0], self.last_pos[1]) |
| draw_px_cross(cr, 5) |
| |
| else: |
| # Draw a filled white rectangle. |
| set_color(cr, palette["WHITE"]) |
| cr.rectangle(-2.0, -2.0, 4.0, 4.0) |
| cr.fill() |
| |
| self.outline.draw(cr) |
| |
| # Draw max radius |
| set_color(cr, palette["BLUE"]) |
| cr.arc(joint_center[0], joint_center[1], l2 + l1, 0, 2.0 * np.pi) |
| with px(cr): |
| cr.stroke() |
| cr.arc(joint_center[0], joint_center[1], l1 - l2, 0, 2.0 * np.pi) |
| with px(cr): |
| cr.stroke() |
| |
| set_color(cr, Color(0.5, 1.0, 1)) |
| |
| if not self.theta_version: |
| theta1, theta2 = to_theta(self.last_pos, |
| self.circular_index_select) |
| x, y = joint_center[0], joint_center[1] |
| cr.move_to(x, y) |
| |
| x += np.cos(theta1) * l1 |
| y += np.sin(theta1) * l1 |
| cr.line_to(x, y) |
| x += np.cos(theta2) * l2 |
| y += np.sin(theta2) * l2 |
| cr.line_to(x, y) |
| with px(cr): |
| cr.stroke() |
| |
| cr.move_to(self.last_pos[0], self.last_pos[1]) |
| set_color(cr, Color(0.0, 1.0, 0.2)) |
| draw_px_cross(cr, 20) |
| |
| if self.theta_version: |
| set_color(cr, Color(0.0, 1.0, 0.2)) |
| cr.move_to(self.last_pos[0], self.last_pos[1]) |
| draw_px_cross(cr, 5) |
| |
| self.outline.draw_theta(cr) |
| |
| set_color(cr, Color(0.0, 0.5, 1.0)) |
| for i in range(len(self.segments)): |
| color = None |
| if i == self.index: |
| continue |
| color = [0, random.random(), 1] |
| random.shuffle(color) |
| set_color(cr, Color(color[0], color[1], color[2])) |
| self.segments[i].DrawTo(cr, self.theta_version) |
| |
| with px(cr): |
| cr.stroke() |
| |
| # Draw current spline in black |
| color = [0, 0, 0] |
| set_color(cr, Color(color[0], color[1], color[2])) |
| self.segments[self.index].DrawTo(cr, self.theta_version) |
| |
| control1 = get_xy(self.segments[self.index].control1) |
| control2 = get_xy(self.segments[self.index].control2) |
| |
| if self.theta_version: |
| control1 = shift_angles(self.segments[self.index].control1) |
| control2 = shift_angles(self.segments[self.index].control2) |
| |
| cr.move_to(control1[0] + 0.02, control1[1]) |
| cr.arc(control1[0], control1[1], 0.02, 0, 2.0 * np.pi) |
| cr.move_to(control2[0] + 0.02, control2[1]) |
| cr.arc(control2[0], control2[1], 0.02, 0, 2.0 * np.pi) |
| |
| with px(cr): |
| cr.stroke() |
| |
| set_color(cr, Color(0.0, 1.0, 0.5)) |
| |
| # Create the plots |
| if self.joint_thetas: |
| if self.joint_points: |
| titles = ["Proximal", "Distal", "Roll joint"] |
| for i in range(len(self.joint_points)): |
| self.axes[i].clear() |
| self.axes[i].plot(self.joint_thetas[0], |
| self.joint_thetas[1][i]) |
| self.axes[i].scatter([self.joint_points[i][0]], |
| [self.joint_points[i][1]], |
| s=10, |
| c="red") |
| self.axes[i].set_title(titles[i]) |
| plt.title("Joint Angle") |
| plt.xlabel("t (0 to 1)") |
| plt.ylabel("theta (rad)") |
| |
| self.fig.canvas.draw() |
| |
| def cur_pt_in_theta(self): |
| if self.theta_version: return self.last_pos |
| return to_theta(self.last_pos, |
| self.circular_index_select, |
| cross_point=-np.pi, |
| die=False) |
| |
| def do_motion(self, event): |
| o_x = event.x |
| o_y = event.y |
| x = event.x - self.window_shape[0] / 2 |
| y = self.window_shape[1] / 2 - event.y |
| scale = self.get_current_scale() |
| event.x = x / scale + self.center[0] |
| event.y = y / scale + self.center[1] |
| |
| segment = self.segments[self.index] |
| self.joint_thetas = segment.joint_thetas() |
| |
| hovered_t = segment.intersection(event) |
| if hovered_t: |
| min_diff = np.inf |
| closest_t = None |
| closest_thetas = None |
| for i in range(len(self.joint_thetas[0])): |
| t = self.joint_thetas[0][i] |
| diff = abs(t - hovered_t) |
| if diff < min_diff: |
| min_diff = diff |
| closest_t = t |
| closest_thetas = [ |
| self.joint_thetas[1][0][i], self.joint_thetas[1][1][i], |
| self.joint_thetas[1][2][i] |
| ] |
| self.joint_points = [(closest_t, closest_theta) |
| for closest_theta in closest_thetas] |
| |
| event.x = o_x |
| event.y = o_y |
| |
| self.redraw() |
| |
| def do_key_press(self, event): |
| keyval = Gdk.keyval_to_lower(event.keyval) |
| print("Gdk.KEY_" + Gdk.keyval_name(keyval)) |
| if keyval == Gdk.KEY_q: |
| print("Found q key and exiting.") |
| quit_main_loop() |
| elif keyval == Gdk.KEY_c: |
| # Increment which arm solution we render |
| self.circular_index_select += 1 |
| print(self.circular_index_select) |
| elif keyval == Gdk.KEY_v: |
| # Decrement which arm solution we render |
| self.circular_index_select -= 1 |
| print(self.circular_index_select) |
| |
| elif keyval == Gdk.KEY_r: |
| self.prev_segment_pt = self.now_segment_pt |
| |
| elif keyval == Gdk.KEY_o: |
| # Only prints current segment |
| print(repr(self.segments[self.index])) |
| elif keyval == Gdk.KEY_g: |
| # Generate theta points. |
| if self.segments: |
| print(repr(self.segments[self.index].ToThetaPoints())) |
| elif keyval == Gdk.KEY_e: |
| best_pt = self.now_segment_pt |
| best_dist = 1e10 |
| for segment in self.segments: |
| d = angle_dist_sqr(segment.start, self.now_segment_pt) |
| if (d < best_dist): |
| best_pt = segment.start |
| best_dist = d |
| d = angle_dist_sqr(segment.end, self.now_segment_pt) |
| if (d < best_dist): |
| best_pt = segment.end |
| best_dist = d |
| self.now_segment_pt = best_pt |
| |
| elif keyval == Gdk.KEY_p: |
| if self.index > 0: |
| self.index -= 1 |
| else: |
| self.index = len(self.segments) - 1 |
| print("Switched to segment:", self.segments[self.index].name) |
| self.segments[self.index].Print(graph_paths.points) |
| |
| elif keyval == Gdk.KEY_n: |
| self.index += 1 |
| self.index = self.index % len(self.segments) |
| print("Switched to segment:", self.segments[self.index].name) |
| self.segments[self.index].Print(graph_paths.points) |
| |
| elif keyval == Gdk.KEY_t: |
| # Toggle between theta and xy renderings |
| if self.theta_version: |
| theta1, theta2 = self.last_pos |
| data = to_xy(theta1, theta2) |
| self.circular_index_select = int( |
| np.floor((theta2 - theta1) / np.pi)) |
| self.last_pos = (data[0], data[1]) |
| else: |
| self.last_pos = self.cur_pt_in_theta() |
| |
| self.theta_version = not self.theta_version |
| self.reinit_extents() |
| |
| elif keyval == Gdk.KEY_z: |
| self.edit_control1 = not self.edit_control1 |
| if self.edit_control1: |
| self.now_segment_pt = self.segments[self.index].control1 |
| else: |
| self.now_segment_pt = self.segments[self.index].control2 |
| if not self.theta_version: |
| data = to_xy(self.now_segment_pt[0], self.now_segment_pt[1]) |
| self.last_pos = (data[0], data[1]) |
| else: |
| self.last_pos = self.now_segment_pt |
| |
| print("self.last_pos: ", self.last_pos, " ci: ", |
| self.circular_index_select) |
| |
| self.redraw() |
| |
| def do_button_press(self, event): |
| last_pos = self.last_pos |
| self.last_pos = (event.x, event.y) |
| pt_theta = self.cur_pt_in_theta() |
| if pt_theta is None: |
| self.last_pos = last_pos |
| return |
| |
| self.now_segment_pt = np.array(shift_angles(pt_theta)) |
| |
| if self.edit_control1: |
| self.segments[self.index].control1 = self.now_segment_pt |
| else: |
| self.segments[self.index].control2 = self.now_segment_pt |
| |
| print('Clicked at theta: np.array([%s, %s])' % |
| (self.now_segment_pt[0], self.now_segment_pt[1])) |
| if not self.theta_version: |
| print( |
| 'Clicked at to_theta_with_circular_index(%.3f, %.3f, circular_index=%d)' |
| % (self.last_pos[0], self.last_pos[1], |
| self.circular_index_select)) |
| |
| self.segments[self.index].Print(graph_paths.points) |
| |
| self.redraw() |
| |
| |
| arm_ui = ArmUi(graph_paths.segments) |
| print('Starting with segment: ', arm_ui.segments[arm_ui.index].name) |
| arm_ui.segments[arm_ui.index].Print(graph_paths.points) |
| basic_window.RunApp() |