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 AnnotationPlugin(Plugin):
def __init__(self, context):
"""
Annotation plugin to create and edit data sets either by manual annotation or automatically, e.g. using
a tracker, and generating larger data sets with data augmentation.
:param context: Parent QT widget
"""
super(AnnotationPlugin, self).__init__(context)
# Widget setup
self.setObjectName('Label Plugin')
self.widget = QWidget()
context.add_widget(self.widget)
self.widget.resize(1800, 1000)
"""left side (current image, grab img button, ...)"""
self.cur_im_widget = ImageWidget(self.widget,
self.roi_callback,
clear_on_click=False)
self.cur_im_widget.setGeometry(QRect(20, 20, 640, 480))
self.grab_img_button = QPushButton(self.widget)
self.grab_img_button.setText("Grab frame")
self.grab_img_button.clicked.connect(self.grab_frame)
self.grab_img_button.setGeometry(QRect(20, 600, 100, 50))
"""right side (selected image, workspace...)"""
self.sel_im_widget = ImageWidget(self.widget,
self.roi_callback,
clear_on_click=True)
self.sel_im_widget.setGeometry(QRect(720, 20, 640, 480))
"""list widgets for images and annotations"""
self.annotation_list_widget = QListWidget(self.widget)
self.annotation_list_widget.setGeometry(QRect(1400, 50, 150, 200))
self.annotation_list_widget.setObjectName("annotation_list_widget")
self.annotation_list_widget.setSelectionMode(
QAbstractItemView.ExtendedSelection)
self.annotation_list_widget.currentItemChanged.connect(
self.select_annotation)
self.image_list_widget = QListWidget(self.widget)
self.image_list_widget.setGeometry(QRect(1550, 50, 250, 500))
self.image_list_widget.setObjectName("image_list_widget")
self.image_list_widget.setSelectionMode(
QAbstractItemView.ExtendedSelection)
self.image_list_widget.currentItemChanged.connect(self.select_image)
self.output_path_edit = QLineEdit(self.widget)
self.output_path_edit.setGeometry(QRect(1400, 20, 300, 30))
self.output_path_edit.setDisabled(True)
self.edit_path_button = QPushButton(self.widget)
self.edit_path_button.setText("set ws")
self.edit_path_button.setGeometry(QRect(1700, 20, 100, 30))
self.edit_path_button.clicked.connect(self.get_workspace)
""" buttons for adding or deleting annotations"""
self.add_annotation_button = QPushButton(self.widget)
self.add_annotation_button.setText("add")
self.add_annotation_button.setGeometry(QRect(1400, 250, 75, 30))
self.add_annotation_button.clicked.connect(self.add_annotation)
self.remove_annotation_button = QPushButton(self.widget)
self.remove_annotation_button.setText("del")
self.remove_annotation_button.setGeometry(QRect(1475, 250, 75, 30))
self.remove_annotation_button.clicked.connect(
self.remove_current_annotation)
"""label combo box, line edit and button for adding labels"""
self.option_selector = QComboBox(self.widget)
self.option_selector.currentIndexChanged.connect(self.class_change)
self.option_selector.setGeometry(1400, 280, 150, 30)
self.label_edit = QLineEdit(self.widget)
self.label_edit.setGeometry(QRect(1400, 310, 100, 30))
self.label_edit.setDisabled(False)
self.edit_label_button = QPushButton(self.widget)
self.edit_label_button.setText("add")
self.edit_label_button.setGeometry(QRect(1500, 310, 50, 30))
self.edit_label_button.clicked.connect(self.add_label)
""" button for image deletion"""
self.remove_image_button = QPushButton(self.widget)
self.remove_image_button.setText("delete image")
self.remove_image_button.setGeometry(QRect(1550, 550, 150, 30))
self.remove_image_button.clicked.connect(self.remove_current_image)
""" export data """
self.gen_data_label = QLabel(self.widget)
self.gen_data_label.setText("Export workspace: ")
self.gen_data_label.setGeometry(QRect(1550, 650, 250, 50))
self.export_ws_button = QPushButton(self.widget)
self.export_ws_button.setText("Export")
self.export_ws_button.setGeometry(QRect(1550, 700, 125, 50))
self.export_ws_button.clicked.connect(self.export_workspace_to_tf)
self.conf_export_button = QPushButton(self.widget)
self.conf_export_button.setText("Configure")
self.conf_export_button.setGeometry(QRect(1675, 700, 125, 50))
self.conf_export_button.clicked.connect(self.set_export_parameters)
""" generate augmented data """
self.gen_data_label = QLabel(self.widget)
self.gen_data_label.setText("Generate augmented dataset:")
self.gen_data_label.setGeometry(QRect(1550, 800, 250, 50))
self.gen_data_button = QPushButton(self.widget)
self.gen_data_button.setText("Generate")
self.gen_data_button.setGeometry(QRect(1550, 850, 125, 50))
self.gen_data_button.clicked.connect(self.generate_augmented_data)
self.gen_data_button = QPushButton(self.widget)
self.gen_data_button.setText("Configure")
self.gen_data_button.setGeometry(QRect(1675, 850, 125, 50))
self.gen_data_button.clicked.connect(
self.set_data_augmentation_parameters)
""" functional stuff"""
self.bridge = CvBridge()
self.sub = None
self.class_id = -1
self.label = ""
self.changes_done = False
self.workspace = None
self.labels = []
self.images_with_annotations = []
self.cur_annotated_image = None
self.cur_annotation_index = -1
self.class_change()
# export parameters
self.default_config_path = None
self.pretrained_graph = None
self.p_test = 0.2
self.batch_size = 12
# data augmentation parameters
self.gen_dir = None
self.num_illuminate = 1
self.num_scale = 1
self.num_blur = 1
def set_data_augmentation_parameters(self):
self.gen_dir = QFileDialog.getExistingDirectory(
self.widget, "Select output directory")
num_illum, ok = QInputDialog.getText(self.widget,
"Illumination changes per image",
"1")
if ok:
try:
self.num_illuminate = int(num_illum)
except ValueError:
pass
num_scale, ok = QInputDialog.getText(self.widget,
"Scaling changes per image", "1")
if ok:
try:
self.num_scale = int(num_scale)
except ValueError:
pass
num_blur, ok = QInputDialog.getText(self.widget,
"Blurring changes per image", "1")
if ok:
try:
self.num_blur = int(num_blur)
except ValueError:
pass
def generate_augmented_data(self):
if self.gen_dir is None:
warning_dialog("Warning", "Set parameters first")
return
data_augmentation.multiply_dataset(self.workspace, self.gen_dir, None,
0, self.num_illuminate,
self.num_scale, self.num_blur, 0)
print("data augmentation done")
def set_export_parameters(self):
""" Set variables for export via QInputDialog and QFileDialog. """
# default config path
config_path = QFileDialog.getOpenFileName(self.widget,
"Select default config")
config_path = str(config_path[0])
file, ext = os.path.splitext(config_path)
if not ext == ".config":
self.default_config_path = None
warning_dialog("warning", "invalid file extension")
return
else:
self.default_config_path = config_path
# pretrained graph
pretrained_graph_dir = QFileDialog.getExistingDirectory(
self.widget, "Select pretrained graph directory")
self.pretrained_graph = pretrained_graph_dir + "/model.ckpt"
# batch size
batch_size, ok = QInputDialog.getText(self.widget, "Set batch size",
"12")
if ok:
try:
self.batch_size = int(batch_size)
except ValueError:
pass
# test percentage
p_test, ok = QInputDialog.getText(self.widget, "Set test percentage",
"0.2")
if ok:
try:
self.p_test = float(p_test)
except ValueError:
pass
def export_workspace_to_tf(self):
""" Export workspace to training formats. """
if self.default_config_path is None or self.pretrained_graph is None:
warning_dialog("Warning",
"define default config and pretrained graph first")
return
tf_utils.export_data_to_tf(self.workspace,
self.images_with_annotations, self.labels,
self.p_test, self.default_config_path,
self.batch_size, self.pretrained_graph,
True)
tf_utils.create_roi_images(self.workspace,
self.images_with_annotations, self.labels)
print("Export done")
def add_label(self):
""" If label doesn't exist yet, add to the list and combo box. """
new_label = str(self.label_edit.text())
if new_label is None or new_label == "":
return
for label in self.labels:
if label[0] == new_label:
warning_dialog("warning",
"label\"" + label[0] + "\" already exists")
return
new_label = list((new_label, 0))
self.labels.append(new_label)
self.option_selector.addItem(new_label[0])
label_file = self.workspace + "/labels.txt"
utils.write_labels(label_file, self.labels)
def add_annotation(self):
""" Add annotation to current image. Bounding box is just a dummy. Use current label. """
if self.class_id == -1 or self.label == "":
warning_dialog("Warning", "select label first")
return
if self.cur_annotated_image is None:
warning_dialog("Warning", "select image first")
return
label = self.option_selector.currentText()
annotation = utils.AnnotationWithBbox(self.class_id, 1.0, 0.5, 0.5, 1,
1)
self.cur_annotated_image.annotation_list.append(annotation)
index = len(self.cur_annotated_image.annotation_list) - 1
self.add_annotation_to_list_widget(index, label, True)
self.changes_done = True
def class_change(self):
""" Called another label is selected in the combo box. Set current label and class id. """
self.label = self.option_selector.currentText()
if self.labels is None or len(self.labels) == 0:
return
self.class_id = [i[0] for i in self.labels].index(self.label)
# num_annotations = self.labels[self.class_id][1]
def select_annotation(self):
"""
Called when an annotation is selected. The corresponding bounding box is drawn thicker,
enables drawing on the image widget.
"""
item = self.annotation_list_widget.currentItem()
if item is None:
self.sel_im_widget.set_active(False)
self.sel_im_widget.clear()
return
text = str(item.text())
index = text.split(":")[0]
self.cur_annotation_index = int(index)
self.show_image_from_workspace()
self.sel_im_widget.set_active(True)
def select_image(self):
"""
Called when an image from the list is selected.
Save annotations, if changes were made. Then show selected image and its annotation list.
"""
if self.changes_done:
utils.save_annotations(self.cur_annotated_image.image_file,
self.cur_annotated_image.annotation_list)
self.changes_done = False
item = self.image_list_widget.currentItem()
if item is None:
return
image_file = self.workspace + str(item.text())
self.cur_annotated_image = None
for img in self.images_with_annotations:
if img.image_file == image_file:
self.cur_annotated_image = img
if self.cur_annotated_image is not None:
self.show_image_from_workspace()
self.set_annotation_list()
def set_annotation_list(self):
""" Set list of annotations for the current image. """
self.annotation_list_widget.clear()
self.cur_annotation_index = -1
for i in range(len(self.cur_annotated_image.annotation_list)):
index = int(self.cur_annotated_image.annotation_list[i].label)
num_labels = len(self.labels)
if index < num_labels:
self.add_annotation_to_list_widget(i, self.labels[index][0])
else:
self.add_annotation_to_list_widget(i, "unknown")
self.show_image_from_workspace()
def add_annotation_to_list_widget(self, index, label, select=False):
"""
Add a QListWidgetItem to the annotation QListWidget
:param index: index in annotation list
:param label: label of annotation
:param select: select new item or not
"""
item = QListWidgetItem()
item.setText(str(index) + ":" + label)
self.annotation_list_widget.addItem(item)
if select:
self.annotation_list_widget.setCurrentItem(item)
self.select_annotation()
def add_image_to_list_widget(self, file_name, select=False):
"""
Add a QListWidgetItem to the image QListWidget
:param file_name: image file name
:param select: select new item or not
"""
item = QListWidgetItem()
item.setText(file_name)
self.image_list_widget.addItem(item)
if select:
self.image_list_widget.setCurrentItem(item)
self.select_image()
def refresh_image_list_widget(self):
""" Clear and fill image list widget """
self.image_list_widget.clear()
for img in self.images_with_annotations:
self.add_image_to_list_widget(
img.image_file.replace(self.workspace, ""))
def remove_current_annotation(self):
""" Remove currently selected annotation """
if self.cur_annotation_index == -1:
warning_dialog("Warning", "no annotation selected")
return
del (self.cur_annotated_image.annotation_list[
self.cur_annotation_index])
self.set_annotation_list()
def remove_current_image(self):
""" Remove currently selected image. Remove ItemWidget and delete files """
if self.cur_annotated_image is None:
warning_dialog("Warning", "no image selected")
return
image_file = self.cur_annotated_image.image_file
label_file = image_file.replace("images",
"labels").replace("jpg", "txt")
self.image_list_widget.removeItemWidget(
self.image_list_widget.currentItem())
self.images_with_annotations.remove(self.cur_annotated_image)
os.remove(image_file)
if os.path.isfile(label_file):
os.remove(label_file)
self.refresh_image_list_widget()
def get_workspace(self):
""" Gets and sets the output directory via a QFileDialog. Save before, if changes were done. """
if self.changes_done:
utils.save_annotations(self.cur_annotated_image.image_file,
self.cur_annotated_image.annotation_list)
self.changes_done = False
self.load_workspace(
QFileDialog.getExistingDirectory(self.widget,
"Select output directory"))
def load_workspace(self, path):
"""
Sets the workspace directory. Checks for missing directories and files,
then loads all images, annotations & label list.
:param path: The path of the directory
"""
if not path:
path = "/tmp"
self.workspace = path
self.output_path_edit.setText(path)
# clear all lists and references
self.labels = []
self.cur_annotation_index = -1
self.cur_annotated_image = None
self.images_with_annotations = []
self.changes_done = False
self.class_id = -1
self.label = ""
# clear combo box and listWidgets
self.option_selector.clear()
self.image_list_widget.clear()
self.annotation_list_widget.clear()
if utils.check_workspace(self.workspace):
image_dir = self.workspace + "/images"
label_file = self.workspace + "/labels.txt"
self.labels = utils.read_labels(label_file)
for label in self.labels:
self.option_selector.addItem(label[0])
for dirname, dirnames, filenames in os.walk(image_dir):
for filename in sorted(filenames):
image_file = dirname + '/' + filename
label_file = image_file.replace(
"images",
"labels").replace(".jpg",
".txt").replace(".png", ".txt")
annotated_image = utils.read_annotated_image(
image_file, label_file)
self.images_with_annotations.append(annotated_image)
self.add_image_to_list_widget("/images/" + filename)
print("workspace loaded successfully")
def grab_frame(self):
""" Grab current frame, save to workspace and show it."""
if self.workspace is None or self.workspace == "":
warning_dialog("Warning", "select workspace first")
return
cv_image = self.cur_im_widget.get_image()
if cv_image is None:
warning_dialog("warning", "no frame to grab")
return
file_name = "/images/{}.jpg".format(
datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S_%f"))
new_image = utils.AnnotatedImage(self.workspace + file_name, [])
self.cur_annotated_image = new_image
self.images_with_annotations.append(new_image)
cv2.imwrite(self.workspace + file_name, cv_image)
self.add_image_to_list_widget(file_name, True)
self.select_image()
# self.sel_im_widget.set_image(cv_image, None, None)
def roi_callback(self):
""" Called when a roi is drawn on an image_widget. Edit annotation if possible. """
if self.class_id == -1 or self.label == "":
warning_dialog("Warning", "select label first")
self.changes_done = True
return
if self.cur_annotation_index < 0:
warning_dialog("Warning", "select annotation first")
self.changes_done = True
return
# set bounding box
x_center, y_center, width, height = self.sel_im_widget.get_normalized_roi(
)
annotation = self.cur_annotated_image.annotation_list[
self.cur_annotation_index]
annotation.bbox = utils.BoundingBox(x_center, y_center, width, height)
# set selected label
self.label = self.option_selector.currentText()
if self.labels is None or len(self.labels) == 0:
return
self.class_id = [i[0] for i in self.labels].index(self.label)
annotation.label = self.class_id
self.set_annotation_list()
self.show_image_from_workspace()
self.sel_im_widget.clear()
self.changes_done = True
def show_image_from_workspace(self):
""" Show current selected image and annotations """
if self.cur_annotated_image is None:
return
img = cv2.imread(self.cur_annotated_image.image_file)
if img is not None:
self.sel_im_widget.set_image(
img, self.cur_annotated_image.annotation_list,
self.cur_annotation_index)
def image_callback(self, msg):
"""
Called when a new sensor_msgs/Image is coming in
:param msg: The image message
"""
try:
cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
except CvBridgeError as e:
rospy.logerr(e)
self.cur_im_widget.set_image(cv_image, None, None)
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)
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
"""
if self.changes_done:
utils.save_annotations(self.cur_annotated_image.image_file,
self.cur_annotated_image.annotation_list)
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("workspace_dir", self.workspace)
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
"""
workspace = None
try:
workspace = instance_settings.value("workspace_dir")
except:
pass
self.load_workspace(workspace)
self.create_subscriber(
str(instance_settings.value("topic_name", "/xtion/rgb/image_raw")))