added qt gui with changeable parameters

remote_control/aruco_estimator.py

@@ -5,6 +5,7 @@ import os
 import time
 import math
 from pyqtgraph.Qt import QtCore, QtGui
+from pyqtgraph.parametertree import Parameter, ParameterTree
 import pyqtgraph as pg
 
 
@@ -13,54 +14,105 @@ from queue import Queue
 
 import aruco
 
+class FPSCounter:
+    def __init__(self):
+        self.fps_start_time = time.time()
+        self.fps_counter = 0
+        self.fps_display_rate = 1.0
+
+    def get_fps(self):
+        self.fps_counter += 1
+        delta_t = time.time() - self.fps_start_time
+        if delta_t > self.fps_display_rate:
+            self.fps_counter = 0
+            self.fps_start_time = time.time()
+        return self.fps_counter / delta_t
 
 class ArucoEstimator:
-    corner_marker_ids = {
-        'a': 0,
-        'b': 1,
-        'c': 2,
-        'd': 3
-    }
-
-    corner_estimates = {
-        'a': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
-        'b': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
-        'c': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
-        'd': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
-    }
-
     def __init__(self, robot_marker_ids=None, use_realsense=True, grid_columns=12, grid_rows=12):
         self.app = QtGui.QApplication([])
 
         ## Create window with GraphicsView widget
-        self.win = pg.GraphicsLayoutWidget()
+        self.win = QtGui.QWidget()
+        self.layout = QtGui.QGridLayout()
+        self.win.setLayout(self.layout)
+
         self.win.keyPressEvent = self.handleKeyPressEvent
-        self.win.show()  ## show widget alone in its own window
         self.win.setWindowTitle('ArucoEstimator')
-        self.view = self.win.addViewBox()
+
+        self.plotwidget = pg.PlotWidget()
+        self.layout.addWidget(self.plotwidget)
 
         ## lock the aspect ratio so pixels are always square
-        self.view.setAspectLocked(True)
+        self.plotwidget.setAspectLocked(True)
+        self.plotwidget.getPlotItem().getAxis('left').hide()
+        self.plotwidget.getPlotItem().getAxis('bottom').hide()
 
         ## Create image item
         self.img = pg.ImageItem(border='w')
         self.img.setLevels([[0, 255], [0, 255], [0, 255]])
         self.img.mouseClickEvent = self.handleMouseEvent
-        self.view.addItem(self.img)
+        self.plotwidget.addItem(self.img)
 
-        self.grid_columns = grid_columns
+        # fps display
-        self.grid_rows = grid_rows
+        self.fps_counter = FPSCounter()
+        self.fps_overlay = pg.TextItem('fps = 0', color=(255, 255, 0), anchor=(0,1))
+        self.plotwidget.addItem(self.fps_overlay)
 
+        self.invert_grayscale = False
+        self.draw_grid = False
+        self.draw_markers = True
+        self.draw_marker_coordinate_system = False
+        self.corner_marker_size = 0.075
+        self.corner_marker_ids = {
+            'a': 0,
+            'b': 1,
+            'c': 2,
+            'd': 3
+        }
+        self.corner_estimates = {
+            'a': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
+            'b': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
+            'c': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
+            'd': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
+        }
+
+        self.robot_marker_size = 0.07
         if robot_marker_ids is None:
             robot_marker_ids = []
         self.robot_marker_ids = robot_marker_ids
         self.robot_marker_estimates = dict([(marker_id, {'t': None, 'x': None, 'y': None, 'angle': None})
                                             for marker_id in self.robot_marker_ids])
 
-        self.draw_grid = False
 
-        self.fps_start_time = time.time()
+        # parameter
-        self.fps_counter = 0
+        params_spec = [
+            {'name': 'Corner marker size', 'type': 'float', 'value': self.corner_marker_size, 'siPrefix': True,
+             'suffix': 'm', 'limits': (0.01, 1), 'step': 0.001},
+            {'name': 'Robot marker size', 'type': 'float', 'value': self.robot_marker_size, 'siPrefix': True,
+             'suffix': 'm', 'limits': (0.01, 1), 'step': 0.001},
+            {'name': 'Draw markers', 'type': 'bool', 'value': self.draw_markers},
+            {'name': 'Draw marker coordinate system', 'type': 'bool', 'value': self.draw_marker_coordinate_system},
+            {'name': 'Invert grayscale', 'type': 'bool', 'value': self.invert_grayscale, 'tip': "Invert grayscale image before marker detection"},
+            {'name': 'Show FPS', 'type': 'bool', 'value': True, 'tip': "Display frames per second counter"},
+            {'name': 'Draw grid', 'type': 'bool', 'value': self.draw_grid, 'tip': "Draw grid spanned by the markers 0 - 3"},
+            {'name': 'Display mode', 'type': 'list', 'values': ['color', 'grayscale'], 'value': 'color', 'tip': "Display mode for the video"}
+            ]
+        self.params = Parameter.create(name='params', type='group', children=params_spec)
+        self.params.param('Invert grayscale').sigValueChanged.connect(lambda _, v: self.__setattr__('invert_grayscale', v))
+        self.params.param('Draw markers').sigValueChanged.connect(lambda _, v: self.__setattr__('draw_markers', v))
+        self.params.param('Draw marker coordinate system').sigValueChanged.connect(lambda _, v: self.__setattr__('draw_marker_coordinate_system', v))
+        self.params.param('Draw grid').sigValueChanged.connect(lambda _, v: self.__setattr__('draw_grid', v))
+        self.params.param('Corner marker size').sigValueChanged.connect(lambda _, v: self.__setattr__('corner_marker_size', v))
+        self.params.param('Robot marker size').sigValueChanged.connect(lambda _, v: self.__setattr__('robot_marker_size', v))
+        self.params.param('Show FPS').sigValueChanged.connect(lambda _, v: self.fps_overlay.show() if v else self.fps_overlay.hide())
+        self.paramtree = ParameterTree()
+        self.paramtree.setParameters(self.params, showTop=False)
+        self.layout.addWidget(self.paramtree)
+        self.win.show()  ## show widget alone in its own window
+
+        self.grid_columns = grid_columns
+        self.grid_rows = grid_rows
 
         self.event_queue = Queue()
 
@@ -113,7 +165,6 @@ class ArucoEstimator:
         self.drag_line_end = None
         self.drag_line_start = None
         self.previous_click = None
-        self.invert_grayscale = False
 
     def compute_clicked_position(self, px, py):
         if self.all_corners_detected():
@@ -176,15 +227,10 @@ class ArucoEstimator:
                 self.drag_line_end = (px, py)
 
     def process_frame(self):
-        draw_markers=True
-        draw_marker_coordinate_system=False
-        #cv2.setMouseCallback('RoboRally', self.mouse_callback)
-        fps_display_rate = 1  # displays the frame rate every 1 second
         if self.pipeline:
             frames = self.pipeline.wait_for_frames()
             color_frame = frames.get_color_frame()
-            # if not color_frame:
+
-            #     continue
             # Convert images to numpy arrays
             color_image = np.asanyarray(color_frame.get_data())
         else:
@@ -206,16 +252,16 @@ class ArucoEstimator:
                 detected_marker_data[marker.id] = {'marker_center': marker.getCenter()}
 
                 if marker.id in self.corner_marker_ids.values():
-                    marker.calculateExtrinsics(0.075, self.camparam)
+                    marker.calculateExtrinsics(self.corner_marker_size, self.camparam)
                 else:
-                    marker.calculateExtrinsics(0.07, self.camparam)
+                    marker.calculateExtrinsics(self.robot_marker_size, self.camparam)
                 detected_marker_data[marker.id]['Rvec'] = marker.Rvec
                 detected_marker_data[marker.id]['Tvec'] = marker.Tvec
 
-                if draw_markers:
+                if self.draw_markers:
                     marker.draw(color_image, np.array([255, 255, 255]), 2, True)
 
-                if draw_marker_coordinate_system:
+                if self.draw_marker_coordinate_system:
                     aruco.CvDrawingUtils.draw3dAxis(color_image, self.camparam, marker.Rvec, marker.Tvec, .1)
 
         # store data
@@ -232,27 +278,22 @@ class ArucoEstimator:
             color_image = cv2.line(color_image, self.drag_line_start, self.drag_line_end, color=(0, 0, 255), thickness=2)
 
         # compute frame rate
-        self.fps_counter += 1
+        self.fps_overlay.setText(f"fps = {self.fps_counter.get_fps():.1f}")
-        delta_t = time.time() - self.fps_start_time
-        if delta_t > fps_display_rate:
-            self.fps_counter = 0
-            self.fps_start_time = time.time()
-        color_image = cv2.putText(color_image, f"fps = {(self.fps_counter / delta_t):.2f}", (10, 25), cv2.FONT_HERSHEY_PLAIN, 2,
-                                  (0, 255, 255),
-                                  thickness=2)
 
         # Show images
         color_image_rgb = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)  # convert to RGB
-        self.img.setImage(np.transpose(np.flipud(color_image_rgb), axes=(1, 0, 2)))
+        display_mode = self.params.param('Display mode').value()
+        if display_mode == 'color':
+            self.img.setImage(np.transpose(np.flipud(color_image_rgb), axes=(1, 0, 2)))
+        elif display_mode == 'grayscale':
+            self.img.setImage(np.transpose(np.flipud(gray)))
 
         QtCore.QTimer.singleShot(1, self.process_frame)
 
     def handleKeyPressEvent(self, ev):
         key = ev.key()
         self.event_queue.put(('key', key))
-        if key == QtCore.Qt.Key_G:
+        if key == QtCore.Qt.Key_Q:
-            self.draw_grid = not self.draw_grid
-        elif key == QtCore.Qt.Key_Q:
             if self.pipeline is not None:
                 # Stop streaming
                 self.pipeline.stop()
@@ -268,6 +309,7 @@ class ArucoEstimator:
                     print("auto exposure ON")
         elif key == QtCore.Qt.Key_I:
             self.invert_grayscale = not self.invert_grayscale
+
     def update_estimate(self, marker_id, pixel_coord_center, rvec, tvec, t_image):
         # update the marker estimate with new data
         if marker_id in self.corner_marker_ids.values():
@@ -464,6 +506,4 @@ if __name__ == "__main__":
     import sys
 
     if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
-        QtGui.QApplication.instance().exec_()
+        QtGui.QApplication.instance().exec_()
-
-    #estimator.run_tracking(draw_markers=True, invert_grayscale=True)