support yolo format annotation file generation

examples/tutorial/apc2016_obj3.json

@@ -1,5 +1,5 @@
 {
-  "version": "4.0.0",
+  "version": "5.2.0.post4",
   "flags": {},
   "shapes": [
     {
@@ -23,6 +23,7 @@
         ]
       ],
       "group_id": null,
+      "description": null,
       "shape_type": "polygon",
       "flags": {}
     },
@@ -71,6 +72,7 @@
         ]
       ],
       "group_id": null,
+      "description": null,
       "shape_type": "polygon",
       "flags": {}
     },
@@ -115,6 +117,7 @@
         ]
       ],
       "group_id": null,
+      "description": null,
       "shape_type": "polygon",
       "flags": {}
     },
@@ -235,8 +238,26 @@
         ]
       ],
       "group_id": null,
+      "description": null,
       "shape_type": "polygon",
       "flags": {}
+    },
+    {
+      "label": "kong_air_dog_squeakair_tennis_ball",
+      "points": [
+        [
+          364.1133004926108,
+          369.756157635468
+        ],
+        [
+          795.6403940886698,
+          557.9334975369458
+        ]
+      ],
+      "group_id": null,
+      "description": "",
+      "shape_type": "rectangle",
+      "flags": {}
     }
   ],
   "imagePath": "apc2016_obj3.json",

examples/tutorial/apc2016_obj3.txt

@@ -0,0 +1 @@
+0 0.4658022228555143 0.3169681893971899 0.11317835769246425 0.15641887671694152

labelme/app.py

@@ -20,7 +20,7 @@
 
 from . import utils
 from labelme.config import get_config
-from labelme.label_file import LabelFile
+from labelme.label_file import LabelFile, GetLabelFileClassFromFormat
 from labelme.label_file import LabelFileError
 from labelme.logger import logger
 from labelme.shape import Shape
@@ -1245,8 +1245,10 @@ def loadFlags(self, flags):
             item.setCheckState(Qt.Checked if flag else Qt.Unchecked)
             self.flag_widget.addItem(item)
 
-    def saveLabels(self, filename):
-        lf = LabelFile()
+    def saveLabels(self, filename, selectedFormat=None):
+        if selectedFormat is None:
+            selectedFormat = "labelme"
+        lf = GetLabelFileClassFromFormat(selectedFormat)()
 
         def format_shape(s):
             data = s.other_data.copy()
@@ -1745,7 +1747,7 @@ def openFile(self, _value=False):
             for fmt in QtGui.QImageReader.supportedImageFormats()
         ]
         filters = self.tr("Image & Label files (%s)") % " ".join(
-            formats + ["*%s" % LabelFile.suffix]
+            formats + ["*%s" % LabelFile.defaultSuffix]
         )
         fileDialog = FileDialogPreview(self)
         fileDialog.setFileMode(FileDialogPreview.ExistingFile)
@@ -1801,20 +1803,31 @@ def saveFile(self, _value=False):
         assert not self.image.isNull(), "cannot save empty image"
         if self.labelFile:
             # DL20180323 - overwrite when in directory
-            self._saveFile(self.labelFile.filename)
+            self._saveFile(self.labelFile.filename, "labelme")
         elif self.output_file:
-            self._saveFile(self.output_file)
+            self._saveFile(self.output_file, "labelme")
             self.close()
         else:
-            self._saveFile(self.saveFileDialog())
+            self._saveFile(*self.saveFileDialog())
 
     def saveFileAs(self, _value=False):
         assert not self.image.isNull(), "cannot save empty image"
-        self._saveFile(self.saveFileDialog())
+        self._saveFile(*self.saveFileDialog())
 
     def saveFileDialog(self):
         caption = self.tr("%s - Choose File") % __appname__
-        filters = self.tr("Label files (*%s)") % LabelFile.suffix
+        filters = self.tr(
+            ";;".join(
+                [
+                    "Label files in %s format (*%s)"
+                    % (oneSuffixFormat, oneSuffix)
+                    for oneSuffix, oneSuffixFormat in zip(
+                        LabelFile.support_suffix_list,
+                        LabelFile.suffix_corresponding_formats,
+                    )
+                ]
+            )
+        )
         if self.output_dir:
             dlg = QtWidgets.QFileDialog(
                 self, caption, self.output_dir, filters
@@ -1823,31 +1836,39 @@ def saveFileDialog(self):
             dlg = QtWidgets.QFileDialog(
                 self, caption, self.currentPath(), filters
             )
-        dlg.setDefaultSuffix(LabelFile.suffix[1:])
+
+        dlg.setDefaultSuffix(LabelFile.support_suffix_list[0][1:])
         dlg.setAcceptMode(QtWidgets.QFileDialog.AcceptSave)
         dlg.setOption(QtWidgets.QFileDialog.DontConfirmOverwrite, False)
         dlg.setOption(QtWidgets.QFileDialog.DontUseNativeDialog, False)
+
         basename = osp.basename(osp.splitext(self.filename)[0])
         if self.output_dir:
             default_labelfile_name = osp.join(
-                self.output_dir, basename + LabelFile.suffix
+                self.output_dir, basename + LabelFile.support_suffix_list[0]
             )
         else:
             default_labelfile_name = osp.join(
-                self.currentPath(), basename + LabelFile.suffix
+                self.currentPath(), basename + LabelFile.support_suffix_list[0]
             )
-        filename = dlg.getSaveFileName(
+        filename, selectedFilter = dlg.getSaveFileName(
             self,
             self.tr("Choose File"),
             default_labelfile_name,
-            self.tr("Label files (*%s)") % LabelFile.suffix,
+            filters,
         )
+        if selectedFilter != "":
+            selectedFormat = selectedFilter.rsplit("Label files in ", 1)[1].split(" format")[0]
+        else:
+            selectedFormat = None
         if isinstance(filename, tuple):
             filename, _ = filename
-        return filename
+        return filename, selectedFormat
 
-    def _saveFile(self, filename):
-        if filename and self.saveLabels(filename):
+    def _saveFile(self, filename, selectedFormat):
+        if filename and self.saveLabels(
+            filename, selectedFormat=selectedFormat
+        ):
             self.addRecentFile(filename)
             self.setClean()
 

labelme/label_file.py

@@ -34,7 +34,9 @@ class LabelFileError(Exception):
 
 class LabelFile(object):
 
-    suffix = ".json"
+    defaultSuffix = ".json"  # by default using labelme format
+    support_suffix_list = [".txt", ".json"]
+    suffix_corresponding_formats = ["yolo", "labelme"]
 
     def __init__(self, filename=None):
         self.shapes = []
@@ -207,4 +209,65 @@ def save(
 
     @staticmethod
     def is_label_file(filename):
-        return osp.splitext(filename)[1].lower() == LabelFile.suffix
+        return osp.splitext(filename)[1].lower() == LabelFile.defaultSuffix
+
+
+class LabelFileYolo(LabelFile):
+    def __init__(self):
+        super(LabelFile, self).__init__()
+
+    def save(
+        self,
+        filename,
+        shapes,
+        imagePath,
+        imageHeight,
+        imageWidth,
+        imageData=None,
+        otherData=None,
+        flags=None,
+    ):
+        try:
+            outputList = []
+            # Note: in yolo annotation, only rectangle bounding box is used.
+            for indexShape, oneShape in enumerate(shapes):
+                centerX = (
+                    (oneShape["points"][0][0] + oneShape["points"][1][0]) / 2 / imageWidth
+                )
+                centerY = (
+                    (oneShape["points"][0][1] + oneShape["points"][1][1]) / 2 / imageHeight
+                )
+                boxWidth = (
+                    abs(oneShape["points"][0][0] - oneShape["points"][1][0]) / imageWidth
+                )
+                boxHeight = (
+                    abs(oneShape["points"][0][1] - oneShape["points"][1][1]) / imageHeight
+                )
+                oneLine = " ".join(
+                    [
+                        oneShape["label"],
+                        str(centerX),
+                        str(centerY),
+                        str(boxWidth),
+                        str(boxHeight),
+                    ]
+                )
+                if indexShape == (len(shapes) - 1):
+                    outputList.append(oneLine)
+                else:
+                    outputList.append(oneLine + "\n")
+            if filename[-4:] != ".txt":
+                filename += ".txt"
+            with open(filename, "w") as outputFile:
+                outputFile.writelines(outputList)
+        except Exception as e:
+            raise LabelFileError(
+                "<Warning> Yolo format only for rectangle annotations!: \n" + e
+            )
+
+
+def GetLabelFileClassFromFormat(selectedFormat):
+    if selectedFormat == "labelme":
+        return LabelFile
+    elif selectedFormat == "yolo":
+        return LabelFileYolo