Add rotated bounding box labeling for YOLO oriented-object-detection

.gitignore

@@ -7,6 +7,7 @@ labelImg.egg-info*
 
 build/
 dist/
+venv/
 
 tags
 cscope*

README.rst

@@ -32,14 +32,17 @@ LabelImg is a graphical image annotation tool.
 It is written in Python and uses Qt for its graphical interface.
 
 Annotations are saved as XML files in PASCAL VOC format, the format used
-by `ImageNet <http://www.image-net.org/>`__.  Besides, it also supports YOLO and CreateML formats.
+by `ImageNet <http://www.image-net.org/>`__.  Besides, it also supports YOLO, RotatedYOLO and CreateML formats.
 
 .. image:: https://raw.githubusercontent.com/tzutalin/labelImg/master/demo/demo3.jpg
      :alt: Demo Image
 
 .. image:: https://raw.githubusercontent.com/tzutalin/labelImg/master/demo/demo.jpg
      :alt: Demo Image
 
+.. image:: /demo/demo6.png
+     :alt: Demo Image
+
 `Watch a demo video <https://youtu.be/p0nR2YsCY_U>`__
 
 Installation
@@ -236,33 +239,41 @@ Annotation visualization
 Hotkeys
 ~~~~~~~
 
-+--------------------+--------------------------------------------+
-| Ctrl + u           | Load all of the images from a directory    |
-+--------------------+--------------------------------------------+
-| Ctrl + r           | Change the default annotation target dir   |
-+--------------------+--------------------------------------------+
-| Ctrl + s           | Save                                       |
-+--------------------+--------------------------------------------+
-| Ctrl + d           | Copy the current label and rect box        |
-+--------------------+--------------------------------------------+
-| Ctrl + Shift + d   | Delete the current image                   |
-+--------------------+--------------------------------------------+
-| Space              | Flag the current image as verified         |
-+--------------------+--------------------------------------------+
-| w                  | Create a rect box                          |
-+--------------------+--------------------------------------------+
-| d                  | Next image                                 |
-+--------------------+--------------------------------------------+
-| a                  | Previous image                             |
-+--------------------+--------------------------------------------+
-| del                | Delete the selected rect box               |
-+--------------------+--------------------------------------------+
-| Ctrl++             | Zoom in                                    |
-+--------------------+--------------------------------------------+
-| Ctrl--             | Zoom out                                   |
-+--------------------+--------------------------------------------+
-| ↑→↓←               | Keyboard arrows to move selected rect box  |
-+--------------------+--------------------------------------------+
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl + u           | Load all of the images from a directory                                                  |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl + r           | Change the default annotation target dir                                                 |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl + s           | Save                                                                                     |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl + d           | Copy the current label and rect box                                                      |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl + Shift + d   | Delete the current image                                                                 |
++--------------------+------------------------------------------------------------------------------------------+
+| Space              | Flag the current image as verified                                                       |
++--------------------+------------------------------------------------------------------------------------------+
+| w                  | Create a rect box                                                                        |
++--------------------+------------------------------------------------------------------------------------------+
+| d                  | Next image                                                                               |
++--------------------+------------------------------------------------------------------------------------------+
+| a                  | Previous image                                                                           |
++--------------------+------------------------------------------------------------------------------------------+
+| del                | Delete the selected rect box                                                             |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl++             | Zoom in                                                                                  |
++--------------------+------------------------------------------------------------------------------------------+
+| Ctrl--             | Zoom out                                                                                 |
++--------------------+------------------------------------------------------------------------------------------+
+| ↑→↓←               | Keyboard arrows to move selected rect box                                                |
++--------------------+------------------------------------------------------------------------------------------+
+| z                  | Counter-clockwise rotation of the selected rect box at a large angle. (RotatedYOLO Only) |
++--------------------+------------------------------------------------------------------------------------------+
+| x                  | Counter-clockwise rotation of the selected rect box at a small angle. (RotatedYOLO Only) |
++--------------------+------------------------------------------------------------------------------------------+
+| c                  | Clockwise rotation of the selected rect box at a small angle. (RotatedYOLO Only)         |
++--------------------+------------------------------------------------------------------------------------------+
+| v                  | Clockwise rotation of the selected rect box at a large angle. (RotatedYOLO Only)         |
++--------------------+------------------------------------------------------------------------------------------+
 
 **Verify Image:**
 

labelImg.py

@@ -45,6 +45,7 @@
 from libs.yolo_io import TXT_EXT
 from libs.create_ml_io import CreateMLReader
 from libs.create_ml_io import JSON_EXT
+from libs.rotated_yolo_io import RotatedYOLOReader
 from libs.ustr import ustr
 from libs.hashableQListWidgetItem import HashableQListWidgetItem
 
@@ -252,6 +253,8 @@ def get_format_meta(format):
                 return '&YOLO', 'format_yolo'
             elif format == LabelFileFormat.CREATE_ML:
                 return '&CreateML', 'format_createml'
+            elif format == LabelFileFormat.ROTATED_YOLO:
+                return '&RotatedYOLO', 'format_rotated_yolo'
 
         save_format = action(get_format_meta(self.label_file_format)[0],
                              self.change_format, 'Ctrl+Y',
@@ -409,6 +412,10 @@ def get_format_meta(format):
             recentFiles=QMenu(get_str('menu_openRecent')),
             labelList=label_menu)
 
+        # Save label to image folder
+        self.save_label_to_image_folder = QAction(get_str('saveLabelToImageFolder'), self)
+        self.save_label_to_image_folder.setCheckable(True)
+        self.save_label_to_image_folder.setChecked(settings.get(SETTING_SAVE_LABEL_TO_IMAGE_FOLDER, False))
         # Auto saving : Enable auto saving if pressing next
         self.auto_saving = QAction(get_str('autoSaveMode'), self)
         self.auto_saving.setCheckable(True)
@@ -431,6 +438,7 @@ def get_format_meta(format):
         add_actions(self.menus.help, (help_default, show_info, show_shortcut))
         add_actions(self.menus.view, (
             self.auto_saving,
+            self.save_label_to_image_folder,
             self.single_class_mode,
             self.display_label_option,
             labels, advanced_mode, None,
@@ -550,6 +558,7 @@ def keyPressEvent(self, event):
 
     # Support Functions #
     def set_format(self, save_format):
+        self.canvas.canDrawRotatedRect = False
         if save_format == FORMAT_PASCALVOC:
             self.actions.save_format.setText(FORMAT_PASCALVOC)
             self.actions.save_format.setIcon(new_icon("format_voc"))
@@ -567,13 +576,22 @@ def set_format(self, save_format):
             self.actions.save_format.setIcon(new_icon("format_createml"))
             self.label_file_format = LabelFileFormat.CREATE_ML
             LabelFile.suffix = JSON_EXT
+
+        elif save_format == FORMAT_ROTATED_YOLO:
+            self.canvas.canDrawRotatedRect = True
+            self.actions.save_format.setText(FORMAT_ROTATED_YOLO)
+            self.actions.save_format.setIcon(new_icon("format_yolo"))
+            self.label_file_format = LabelFileFormat.ROTATED_YOLO
+            LabelFile.suffix = TXT_EXT
 
     def change_format(self):
         if self.label_file_format == LabelFileFormat.PASCAL_VOC:
             self.set_format(FORMAT_YOLO)
         elif self.label_file_format == LabelFileFormat.YOLO:
             self.set_format(FORMAT_CREATEML)
         elif self.label_file_format == LabelFileFormat.CREATE_ML:
+            self.set_format(FORMAT_ROTATED_YOLO)
+        elif self.label_file_format == LabelFileFormat.ROTATED_YOLO:
             self.set_format(FORMAT_PASCALVOC)
         else:
             raise ValueError('Unknown label file format.')
@@ -724,7 +742,7 @@ def toggle_draw_mode(self, edit=True):
     def set_create_mode(self):
         assert self.advanced()
         self.toggle_draw_mode(False)
-
+    
     def set_edit_mode(self):
         assert self.advanced()
         self.toggle_draw_mode(True)
@@ -837,7 +855,7 @@ def remove_label(self, shape):
 
     def load_labels(self, shapes):
         s = []
-        for label, points, line_color, fill_color, difficult in shapes:
+        for label, points, line_color, fill_color, difficult, direction in shapes:
             shape = Shape(label=label)
             for x, y in points:
 
@@ -848,6 +866,7 @@ def load_labels(self, shapes):
 
                 shape.add_point(QPointF(x, y))
             shape.difficult = difficult
+            shape.direction = direction
             shape.close()
             s.append(shape)
 
@@ -888,7 +907,10 @@ def format_shape(s):
                         fill_color=s.fill_color.getRgb(),
                         points=[(p.x(), p.y()) for p in s.points],
                         # add chris
-                        difficult=s.difficult)
+                        difficult=s.difficult,
+                        direction=s.direction,
+                        center=s.center,
+            )
 
         shapes = [format_shape(shape) for shape in self.canvas.shapes]
         # Can add different annotation formats here
@@ -908,6 +930,11 @@ def format_shape(s):
                     annotation_file_path += JSON_EXT
                 self.label_file.save_create_ml_format(annotation_file_path, shapes, self.file_path, self.image_data,
                                                       self.label_hist, self.line_color.getRgb(), self.fill_color.getRgb())
+            elif self.label_file_format == LabelFileFormat.ROTATED_YOLO:
+                if annotation_file_path[-4:].lower() != ".txt":
+                    annotation_file_path += TXT_EXT
+                self.label_file.save_rotated_yolo_format(annotation_file_path, shapes, self.file_path, self.image_data, self.label_hist,
+                                                 self.line_color.getRgb(), self.fill_color.getRgb())
             else:
                 self.label_file.save(annotation_file_path, shapes, self.file_path, self.image_data,
                                      self.line_color.getRgb(), self.fill_color.getRgb())
@@ -1099,6 +1126,9 @@ def load_file(self, file_path=None):
         # Make sure that filePath is a regular python string, rather than QString
         file_path = ustr(file_path)
 
+        if self.save_label_to_image_folder.isChecked():
+            self.default_save_dir = os.path.dirname(file_path)
+
         # Fix bug: An  index error after select a directory when open a new file.
         unicode_file_path = ustr(file_path)
         unicode_file_path = os.path.abspath(unicode_file_path)
@@ -1178,33 +1208,42 @@ def counter_str(self):
         return '[{} / {}]'.format(self.cur_img_idx + 1, self.img_count)
 
     def show_bounding_box_from_annotation_file(self, file_path):
-        if self.default_save_dir is not None:
-            basename = os.path.basename(os.path.splitext(file_path)[0])
-            xml_path = os.path.join(self.default_save_dir, basename + XML_EXT)
-            txt_path = os.path.join(self.default_save_dir, basename + TXT_EXT)
-            json_path = os.path.join(self.default_save_dir, basename + JSON_EXT)
-
-            """Annotation file priority:
-            PascalXML > YOLO
-            """
-            if os.path.isfile(xml_path):
-                self.load_pascal_xml_by_filename(xml_path)
-            elif os.path.isfile(txt_path):
-                self.load_yolo_txt_by_filename(txt_path)
-            elif os.path.isfile(json_path):
-                self.load_create_ml_json_by_filename(json_path, file_path)
+        if  file_path is not None:
+            if self.default_save_dir is not None:
+                basename = os.path.basename(os.path.splitext(file_path)[0])
+                xml_path = os.path.join(self.default_save_dir, basename + XML_EXT)
+                txt_path = os.path.join(self.default_save_dir, basename + TXT_EXT)
+                json_path = os.path.join(self.default_save_dir, basename + JSON_EXT)
+
+                """Annotation file priority:
+                PascalXML > YOLO
+                """
+                if os.path.isfile(xml_path):
+                    self.load_pascal_xml_by_filename(xml_path)
+                elif os.path.isfile(txt_path):
+                    self.check_txt_label_type(txt_path)
+                    if self.label_file_format == LabelFileFormat.ROTATED_YOLO:
+                        self.load_rotated_yolo_txt_by_filename(txt_path)
+                    else:
+                        self.load_yolo_txt_by_filename(txt_path)
+                elif os.path.isfile(json_path):
+                    self.load_create_ml_json_by_filename(json_path, file_path)
 
-        else:
-            xml_path = os.path.splitext(file_path)[0] + XML_EXT
-            txt_path = os.path.splitext(file_path)[0] + TXT_EXT
-            json_path = os.path.splitext(file_path)[0] + JSON_EXT
-
-            if os.path.isfile(xml_path):
-                self.load_pascal_xml_by_filename(xml_path)
-            elif os.path.isfile(txt_path):
-                self.load_yolo_txt_by_filename(txt_path)
-            elif os.path.isfile(json_path):
-                self.load_create_ml_json_by_filename(json_path, file_path)
+            else:
+                xml_path = os.path.splitext(file_path)[0] + XML_EXT
+                txt_path = os.path.splitext(file_path)[0] + TXT_EXT
+                json_path = os.path.splitext(file_path)[0] + JSON_EXT
+
+                if os.path.isfile(xml_path):
+                    self.load_pascal_xml_by_filename(xml_path)
+                elif os.path.isfile(txt_path):
+                    self.check_txt_label_type(txt_path)
+                    if self.label_file_format == LabelFileFormat.ROTATED_YOLO:
+                        self.load_rotated_yolo_txt_by_filename(txt_path)
+                    else:
+                        self.load_yolo_txt_by_filename(txt_path)
+                elif os.path.isfile(json_path):
+                    self.load_create_ml_json_by_filename(json_path, file_path)
 
 
     def resizeEvent(self, event):
@@ -1270,6 +1309,7 @@ def closeEvent(self, event):
             settings[SETTING_LAST_OPEN_DIR] = ''
 
         settings[SETTING_AUTO_SAVE] = self.auto_saving.isChecked()
+        settings[SETTING_SAVE_LABEL_TO_IMAGE_FOLDER] = self.save_label_to_image_folder.isChecked()
         settings[SETTING_SINGLE_CLASS] = self.single_class_mode.isChecked()
         settings[SETTING_PAINT_LABEL] = self.display_label_option.isChecked()
         settings[SETTING_DRAW_SQUARE] = self.draw_squares_option.isChecked()
@@ -1638,10 +1678,29 @@ def load_yolo_txt_by_filename(self, txt_path):
         self.set_format(FORMAT_YOLO)
         t_yolo_parse_reader = YoloReader(txt_path, self.image)
         shapes = t_yolo_parse_reader.get_shapes()
-        print(shapes)
         self.load_labels(shapes)
         self.canvas.verified = t_yolo_parse_reader.verified
 
+    def load_rotated_yolo_txt_by_filename(self, txt_path):
+        if self.file_path is None:
+            return
+        if os.path.isfile(txt_path) is False:
+            return
+
+        self.set_format(FORMAT_ROTATED_YOLO)
+        t_yolo_parse_reader = RotatedYOLOReader(txt_path)
+        shapes = t_yolo_parse_reader.get_shapes()
+        self.load_labels(shapes)
+        self.canvas.verified = t_yolo_parse_reader.verified
+
+    def check_txt_label_type(self, txt_path):
+        with open(txt_path, 'r') as f:
+            line = f.readline()
+            if len(line.strip().split()) == 5:
+                self.set_format(FORMAT_YOLO)
+            else:
+                self.set_format(FORMAT_ROTATED_YOLO)
+
     def load_create_ml_json_by_filename(self, json_path, file_path):
         if self.file_path is None:
             return