class ManualPlugin(Plugin): def __init__(self, context): """ ManualPlugin class that performs a manual recognition based on a request :param context: QT context, aka parent """ super(ManualPlugin, self).__init__(context) # Widget setup self.setObjectName('Manual Plugin') self._widget = QWidget() context.add_widget(self._widget) # Layout and attach to widget layout = QVBoxLayout() self._widget.setLayout(layout) self._image_widget = ImageWidget(self._widget, self.image_roi_callback) layout.addWidget(self._image_widget) # Input field grid_layout = QGridLayout() layout.addLayout(grid_layout) self._labels_edit = QLineEdit() self._labels_edit.setDisabled(True) grid_layout.addWidget(self._labels_edit, 2, 2) self._edit_labels_button = QPushButton("Edit labels") self._edit_labels_button.clicked.connect(self._get_labels) grid_layout.addWidget(self._edit_labels_button, 2, 1) self._done_recognizing_button = QPushButton("Done recognizing..") self._done_recognizing_button.clicked.connect(self._done_recognizing) self._done_recognizing_button.setDisabled(True) grid_layout.addWidget(self._done_recognizing_button, 3, 2) # Bridge for opencv conversion self.bridge = CvBridge() # Set service to None self._srv = None self._srv_name = None self._response = RecognizeResponse() self._recognizing = False def _get_labels(self): """ Gets and sets the labels """ text, ok = QInputDialog.getText(self._widget, 'Text Input Dialog', 'Type labels semicolon separated, e.g. banana;apple:', QLineEdit.Normal, ";".join(self.labels)) if ok: # Sanitize to alphanumeric, exclude spaces labels = set([_sanitize(label) for label in str(text).split(";") if _sanitize(label)]) self._set_labels(labels) def _set_labels(self, labels): """ Sets the labels :param labels: label string array """ if not labels: labels = [] self.labels = labels self._labels_edit.setText("%s" % labels) def _done_recognizing(self): self._image_widget.clear() self._recognizing = False def recognize_srv_callback(self, req): """ Method callback for the Recognize.srv :param req: The service request """ self._response.recognitions = [] self._recognizing = True try: cv_image = self.bridge.imgmsg_to_cv2(req.image, "bgr8") except CvBridgeError as e: rospy.logerr(e) self._image_widget.set_image(cv_image) self._done_recognizing_button.setDisabled(False) timeout = 60.0 # Maximum of 60 seconds future = rospy.Time.now() + rospy.Duration(timeout) rospy.loginfo("Waiting for manual recognition, maximum of %d seconds", timeout) while not rospy.is_shutdown() and self._recognizing: if rospy.Time.now() > future: raise rospy.ServiceException("Timeout of %d seconds exceeded .." % timeout) rospy.sleep(rospy.Duration(0.1)) self._done_recognizing_button.setDisabled(True) return self._response def image_roi_callback(self, roi_image): """ Callback triggered when the user has drawn an ROI on the image :param roi_image: The opencv image in the ROI """ if not self.labels: warning_dialog("No labels specified!", "Please first specify some labels using the 'Edit labels' button") return height, width = roi_image.shape[:2] option = option_dialog("Label", self.labels) if option: self._image_widget.add_detection(0, 0, width, height, option) self._stage_recognition(self._image_widget.get_roi(), option) def _stage_recognition(self, roi, label): """ Stage a manual recognition :param roi: ROI :param label: The label """ x, y, width, height = roi r = Recognition(roi=RegionOfInterest(x_offset=x, y_offset=y, width=width, height=height)) r.categorical_distribution.probabilities = [CategoryProbability(label=label, probability=1.0)] r.categorical_distribution.unknown_probability = 0.0 self._response.recognitions.append(r) def trigger_configuration(self): """ Callback when the configuration button is clicked """ srv_name, ok = QInputDialog.getText(self._widget, "Select service name", "Service name") if ok: self._create_service_server(srv_name) def _create_service_server(self, srv_name): """ Method that creates a service server for a Recognize.srv :param srv_name: """ if self._srv: self._srv.shutdown() if srv_name: rospy.loginfo("Creating service '%s'" % srv_name) self._srv_name = srv_name self._srv = rospy.Service(srv_name, Recognize, self.recognize_srv_callback) def shutdown_plugin(self): """ Callback function when shutdown is requested """ pass def save_settings(self, plugin_settings, instance_settings): """ Callback function on shutdown to store the local plugin variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ instance_settings.set_value("labels", self.labels) if self._srv: instance_settings.set_value("srv_name", self._srv_name) def restore_settings(self, plugin_settings, instance_settings): """ Callback function fired on load of the plugin that allows to restore saved variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ labels = None try: labels = instance_settings.value("labels") except: pass self._set_labels(labels) self._create_service_server(str(instance_settings.value("srv_name", "/my_recognition_service")))
class TestPlugin(Plugin): def __init__(self, context): """ TestPlugin class to evaluate the image_recognition_msgs interfaces :param context: QT context, aka parent """ super(TestPlugin, self).__init__(context) # Widget setup self.setObjectName('Test Plugin') self._widget = QWidget() context.add_widget(self._widget) # Layout and attach to widget layout = QVBoxLayout() self._widget.setLayout(layout) self._image_widget = ImageWidget(self._widget, self.image_roi_callback, clear_on_click=True) layout.addWidget(self._image_widget) # Input field grid_layout = QGridLayout() layout.addLayout(grid_layout) self._info = QLineEdit() self._info.setDisabled(True) self._info.setText( "Draw a rectangle on the screen to perform recognition of that ROI" ) layout.addWidget(self._info) # Bridge for opencv conversion self.bridge = CvBridge() # Set subscriber and service to None self._sub = None self._srv = None def recognize_srv_call(self, roi_image): """ Method that calls the Recognize.srv :param roi_image: Selected roi_image by the user """ try: result = self._srv( image=self.bridge.cv2_to_imgmsg(roi_image, "bgr8")) except Exception as e: warning_dialog("Service Exception", str(e)) return print result for r in result.recognitions: text_array = [] best = CategoryProbability( label="unknown", probability=r.categorical_distribution.unknown_probability) for p in r.categorical_distribution.probabilities: text_array.append("%s: %.2f" % (p.label, p.probability)) if p.probability > best.probability: best = p self._image_widget.add_detection(r.roi.x_offset, r.roi.y_offset, r.roi.width, r.roi.height, best.label) if text_array: option_dialog( "Classification results (Unknown probability=%.2f)" % r.categorical_distribution.unknown_probability, text_array) # Show all results in a dropdown def get_face_properties_srv_call(self, roi_image): """ Method that calls the GetFaceProperties.srv :param roi_image: Selected roi_image by the user """ try: result = self._srv(face_image_array=[ self.bridge.cv2_to_imgmsg(roi_image, "bgr8") ]) except Exception as e: warning_dialog("Service Exception", str(e)) return msg = "" for properties in result.properties_array: msg += "- FaceProperties(gender=%s, age=%s, glasses=%s, mood=%s)" % \ ("male" if properties.gender == FaceProperties.MALE else "female", properties.age, "false" if properties.glasses == 0 else "true", properties.mood) info_dialog("Face Properties array", msg) def image_roi_callback(self, roi_image): """ Callback triggered when the user has drawn an ROI on the image :param roi_image: The opencv image in the ROI """ if self._srv is None: warning_dialog( "No service specified!", "Please first specify a service via the options button (top-right gear wheel)" ) return if self._srv.service_class == Recognize: self.recognize_srv_call(roi_image) elif self._srv.service_class == GetFaceProperties: self.get_face_properties_srv_call(roi_image) else: warning_dialog("Unknown service class", "Service class is unkown!") def _image_callback(self, msg): """ Sensor_msgs/Image callback :param msg: The image message """ try: cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8") except CvBridgeError as e: rospy.logerr(e) return self._image_widget.set_image(cv_image) def trigger_configuration(self): """ Callback when the configuration button is clicked """ topic_name, ok = QInputDialog.getItem( self._widget, "Select topic name", "Topic name", rostopic.find_by_type('sensor_msgs/Image')) if ok: self._create_subscriber(topic_name) available_rosservices = [] for s in rosservice.get_service_list(): try: if rosservice.get_service_type(s) in _SUPPORTED_SERVICES: available_rosservices.append(s) except: pass srv_name, ok = QInputDialog.getItem(self._widget, "Select service name", "Service name", available_rosservices) if ok: self._create_service_client(srv_name) def _create_subscriber(self, topic_name): """ Method that creates a subscriber to a sensor_msgs/Image topic :param topic_name: The topic_name """ if self._sub: self._sub.unregister() self._sub = rospy.Subscriber(topic_name, Image, self._image_callback) rospy.loginfo("Listening to %s -- spinning .." % self._sub.name) self._widget.setWindowTitle("Test plugin, listening to (%s)" % self._sub.name) def _create_service_client(self, srv_name): """ Method that creates a client service proxy to call either the GetFaceProperties.srv or the Recognize.srv :param srv_name: """ if self._srv: self._srv.close() if srv_name in rosservice.get_service_list(): rospy.loginfo("Creating proxy for service '%s'" % srv_name) self._srv = rospy.ServiceProxy( srv_name, rosservice.get_service_class_by_name(srv_name)) def shutdown_plugin(self): """ Callback function when shutdown is requested """ pass def save_settings(self, plugin_settings, instance_settings): """ Callback function on shutdown to store the local plugin variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ if self._sub: instance_settings.set_value("topic_name", self._sub.name) def restore_settings(self, plugin_settings, instance_settings): """ Callback function fired on load of the plugin that allows to restore saved variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ self._create_subscriber( str(instance_settings.value("topic_name", "/usb_cam/image_raw"))) self._create_service_client( str( instance_settings.value("service_name", "/image_recognition/my_service")))
class AnnotationPlugin(Plugin): def __init__(self, context): """ Annotation plugin to create data sets or test the Annotate.srv service :param context: Parent QT widget """ super(AnnotationPlugin, self).__init__(context) # Widget setup self.setObjectName('Label Plugin') self._widget = QWidget() context.add_widget(self._widget) # Layout and attach to widget layout = QVBoxLayout() self._widget.setLayout(layout) self._image_widget = ImageWidget(self._widget, self.image_roi_callback, clear_on_click=True) layout.addWidget(self._image_widget) # Input field grid_layout = QGridLayout() layout.addLayout(grid_layout) self._edit_path_button = QPushButton("Edit path") self._edit_path_button.clicked.connect(self._get_output_directory) grid_layout.addWidget(self._edit_path_button, 1, 1) self._output_path_edit = QLineEdit() self._output_path_edit.setDisabled(True) grid_layout.addWidget(self._output_path_edit, 1, 2) self._labels_edit = QLineEdit() self._labels_edit.setDisabled(True) grid_layout.addWidget(self._labels_edit, 2, 2) self._edit_labels_button = QPushButton("Edit labels") self._edit_labels_button.clicked.connect(self._get_labels) grid_layout.addWidget(self._edit_labels_button, 2, 1) self._save_button = QPushButton("Annotate again!") self._save_button.clicked.connect(self.annotate_again_clicked) grid_layout.addWidget(self._save_button, 2, 3) # Bridge for opencv conversion self.bridge = CvBridge() # Set subscriber to None self._sub = None self._srv = None self.labels = [] self.label = "" self.output_directory = "" def image_roi_callback(self, roi_image): """ Callback from the image widget when the user has selected a ROI :param roi_image: The opencv image of the ROI """ if not self.labels: warning_dialog( "No labels specified!", "Please first specify some labels using the 'Edit labels' button" ) return height, width = roi_image.shape[:2] option = option_dialog("Label", self.labels) if option: self.label = option self._image_widget.add_detection(0, 0, width, height, option) self.annotate(roi_image) def annotate_again_clicked(self): """ Triggered when button clicked """ roi_image = self._image_widget.get_roi_image() if roi_image is not None: self.annotate(roi_image) def annotate(self, roi_image): """ Create an annotation :param roi_image: The image we want to annotate """ self.annotate_srv(roi_image) self.store_image(roi_image) def annotate_srv(self, roi_image): """ Call the selected Annotate.srv :param roi_image: The full opencv image we want to annotate """ if roi_image is not None and self.label is not None and self._srv is not None: height, width = roi_image.shape[:2] try: self._srv(image=self.bridge.cv2_to_imgmsg(roi_image, "bgr8"), annotations=[ Annotation(label=self.label, roi=RegionOfInterest(x_offset=0, y_offset=0, width=width, height=height)) ]) except Exception as e: warning_dialog("Service Exception", str(e)) def _create_service_client(self, srv_name): """ Create a service client proxy :param srv_name: Name of the service """ if self._srv: self._srv.close() if srv_name in rosservice.get_service_list(): rospy.loginfo("Creating proxy for service '%s'" % srv_name) self._srv = rospy.ServiceProxy( srv_name, rosservice.get_service_class_by_name(srv_name)) def store_image(self, roi_image): """ Store the image :param roi_image: Image we would like to store """ if roi_image is not None and self.label is not None and self.output_directory is not None: image_writer.write_annotated(self.output_directory, roi_image, self.label, True) def _get_output_directory(self): """ Gets and sets the output directory via a QFileDialog """ self._set_output_directory( QFileDialog.getExistingDirectory(self._widget, "Select output directory")) def _set_output_directory(self, path): """ Sets the output directory :param path: The path of the directory """ if not path: path = "/tmp" self.output_directory = path self._output_path_edit.setText("Saving images to %s" % path) def _get_labels(self): """ Gets and sets the labels """ text, ok = QInputDialog.getText( self._widget, 'Text Input Dialog', 'Type labels semicolon separated, e.g. banana;apple:', QLineEdit.Normal, ";".join(self.labels)) if ok: labels = set([ _sanitize(label) for label in str(text).split(";") if _sanitize(label) ]) # Sanitize to alphanumeric, exclude spaces self._set_labels(labels) def _set_labels(self, labels): """ Sets the labels :param labels: label string array """ if not labels: labels = [] self.labels = labels self._labels_edit.setText("%s" % labels) def _image_callback(self, msg): """ Called when a new sensor_msgs/Image is coming in :param msg: The image messaeg """ try: cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8") except CvBridgeError as e: rospy.logerr(e) self._image_widget.set_image(cv_image) def trigger_configuration(self): """ Callback when the configuration button is clicked """ topic_name, ok = QInputDialog.getItem( self._widget, "Select topic name", "Topic name", rostopic.find_by_type('sensor_msgs/Image')) if ok: self._create_subscriber(topic_name) available_rosservices = [] for s in rosservice.get_service_list(): try: if rosservice.get_service_type(s) in _SUPPORTED_SERVICES: available_rosservices.append(s) except: pass srv_name, ok = QInputDialog.getItem(self._widget, "Select service name", "Service name", available_rosservices) if ok: self._create_service_client(srv_name) def _create_subscriber(self, topic_name): """ Method that creates a subscriber to a sensor_msgs/Image topic :param topic_name: The topic_name """ if self._sub: self._sub.unregister() self._sub = rospy.Subscriber(topic_name, Image, self._image_callback) rospy.loginfo("Listening to %s -- spinning .." % self._sub.name) self._widget.setWindowTitle("Label plugin, listening to (%s)" % self._sub.name) def shutdown_plugin(self): """ Callback function when shutdown is requested """ pass def save_settings(self, plugin_settings, instance_settings): """ Callback function on shutdown to store the local plugin variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ instance_settings.set_value("output_directory", self.output_directory) instance_settings.set_value("labels", self.labels) if self._sub: instance_settings.set_value("topic_name", self._sub.name) def restore_settings(self, plugin_settings, instance_settings): """ Callback function fired on load of the plugin that allows to restore saved variables :param plugin_settings: Plugin settings :param instance_settings: Settings of this instance """ path = None try: path = instance_settings.value("output_directory") except: pass self._set_output_directory(path) labels = None try: labels = instance_settings.value("labels") except: pass self._set_labels(labels) self._create_subscriber( str(instance_settings.value("topic_name", "/usb_cam/image_raw"))) self._create_service_client( str( instance_settings.value("service_name", "/image_recognition/my_service")))