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")))