Adds create-ml format support (#651)

* adds createMl reader & writer class

* adds getFormatMeta function to support more than two save_format

* adds CreateML read & write support

* adds format CreateML icon

* fixes negative height/width

* removes type hints

* fixes coordinate calculation

* adds unit test

* removes typehint

labelImg.py

@@ -44,11 +44,14 @@ from libs.pascal_voc_io import PascalVocReader
 from libs.pascal_voc_io import XML_EXT
 from libs.yolo_io import YoloReader
 from libs.yolo_io import TXT_EXT
+from libs.create_ml_io import CreateMLReader
+from libs.create_ml_io import JSON_EXT
 from libs.ustr import ustr
 from libs.hashableQListWidgetItem import HashableQListWidgetItem
 
 __appname__ = 'labelImg'
 
+
 class WindowMixin(object):
 
     def menu(self, title, actions=None):
@@ -231,10 +234,20 @@ class MainWindow(QMainWindow, WindowMixin):
         save = action(getStr('save'), self.saveFile,
                       'Ctrl+S', 'save', getStr('saveDetail'), enabled=False)
 
-        isUsingPascalVoc = self.labelFileFormat == LabelFileFormat.PASCAL_VOC
-        save_format = action('&PascalVOC' if isUsingPascalVoc else '&YOLO',
+        def getFormatMeta(format):
+            """
+            returns a tuple containing (title, icon_name) of the selected format
+            """
+            if format == LabelFileFormat.PASCAL_VOC:
+                return ('&PascalVOC', 'format_voc')
+            elif format == LabelFileFormat.YOLO:
+                return ('&YOLO', 'format_yolo')
+            elif format == LabelFileFormat.CREATE_ML:
+                return ('&CreateML', 'format_createml')
+
+        save_format = action(getFormatMeta(self.labelFileFormat)[0],
                              self.change_format, 'Ctrl+',
-                             'format_voc' if isUsingPascalVoc else 'format_yolo',
+                             getFormatMeta(self.labelFileFormat)[1],
                              getStr('changeSaveFormat'), enabled=True)
 
         saveAs = action(getStr('saveAs'), self.saveFileAs,
@@ -515,10 +528,18 @@ class MainWindow(QMainWindow, WindowMixin):
             self.labelFileFormat = LabelFileFormat.YOLO
             LabelFile.suffix = TXT_EXT
 
+        elif save_format == FORMAT_CREATEML:
+            self.actions.save_format.setText(FORMAT_CREATEML)
+            self.actions.save_format.setIcon(newIcon("format_createml"))
+            self.labelFileFormat = LabelFileFormat.CREATE_ML
+            LabelFile.suffix = JSON_EXT
+
     def change_format(self):
         if self.labelFileFormat == LabelFileFormat.PASCAL_VOC:
             self.set_format(FORMAT_YOLO)
         elif self.labelFileFormat == LabelFileFormat.YOLO:
+            self.set_format(FORMAT_CREATEML)
+        elif self.labelFileFormat == LabelFileFormat.CREATE_ML:
             self.set_format(FORMAT_PASCALVOC)
         else:
             raise ValueError('Unknown label file format.')
@@ -835,6 +856,11 @@ class MainWindow(QMainWindow, WindowMixin):
                     annotationFilePath += TXT_EXT
                 self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
                                               self.lineColor.getRgb(), self.fillColor.getRgb())
+            elif self.labelFileFormat == LabelFileFormat.CREATE_ML:
+                if annotationFilePath[-5:].lower() != ".json":
+                    annotationFilePath += JSON_EXT
+                self.labelFile.saveCreateMLFormat(annotationFilePath, shapes, self.filePath, self.imageData,
+                                                  self.labelHist, self.lineColor.getRgb(), self.fillColor.getRgb())
             else:
                 self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
                                     self.lineColor.getRgb(), self.fillColor.getRgb())
@@ -1084,8 +1110,12 @@ class MainWindow(QMainWindow, WindowMixin):
         if self.defaultSaveDir is not None:
             basename = os.path.basename(
                 os.path.splitext(filePath)[0])
+
+            filedir = filePath.split(basename)[0].split("/")[-2:-1][0]
+
             xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT)
             txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT)
+            jsonPath = os.path.join(self.defaultSaveDir, filedir + JSON_EXT)
 
             """Annotation file priority:
             PascalXML > YOLO
@@ -1094,6 +1124,9 @@ class MainWindow(QMainWindow, WindowMixin):
                 self.loadPascalXMLByFilename(xmlPath)
             elif os.path.isfile(txtPath):
                 self.loadYOLOTXTByFilename(txtPath)
+            elif os.path.isfile(jsonPath):
+                self.loadCreateMLJSONByFilename(jsonPath, filePath)
+
         else:
             xmlPath = os.path.splitext(filePath)[0] + XML_EXT
             txtPath = os.path.splitext(filePath)[0] + TXT_EXT
@@ -1502,6 +1535,19 @@ class MainWindow(QMainWindow, WindowMixin):
         self.loadLabels(shapes)
         self.canvas.verified = tYoloParseReader.verified
 
+    def loadCreateMLJSONByFilename(self, jsonPath, filePath):
+        if self.filePath is None:
+            return
+        if os.path.isfile(jsonPath) is False:
+            return
+
+        self.set_format(FORMAT_CREATEML)
+
+        crmlParseReader = CreateMLReader(jsonPath, filePath)
+        shapes = crmlParseReader.get_shapes()
+        self.loadLabels(shapes)
+        self.canvas.verified = crmlParseReader.verified
+
     def copyPreviousBoundingBoxes(self):
         currIndex = self.mImgList.index(self.filePath)
         if currIndex - 1 >= 0:

libs/constants.py

@@ -14,6 +14,7 @@ SETTING_AUTO_SAVE = 'autosave'
 SETTING_SINGLE_CLASS = 'singleclass'
 FORMAT_PASCALVOC='PascalVOC'
 FORMAT_YOLO='YOLO'
+FORMAT_CREATEML='CreateML'
 SETTING_DRAW_SQUARE = 'draw/square'
 SETTING_LABEL_FILE_FORMAT= 'labelFileFormat'
 DEFAULT_ENCODING = 'utf-8'

libs/create_ml_io.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python
+# -*- coding: utf8 -*-
+import json
+from pathlib import Path
+
+from libs.constants import DEFAULT_ENCODING
+import os
+
+JSON_EXT = '.json'
+ENCODE_METHOD = DEFAULT_ENCODING
+
+
+class CreateMLWriter:
+    def __init__(self, foldername, filename, imgsize, shapes, outputfile, databasesrc='Unknown', localimgpath=None):
+        self.foldername = foldername
+        self.filename = filename
+        self.databasesrc = databasesrc
+        self.imgsize = imgsize
+        self.boxlist = []
+        self.localimgpath = localimgpath
+        self.verified = False
+        self.shapes = shapes
+        self.outputfile = outputfile
+
+    def write(self):
+        if os.path.isfile(self.outputfile):
+            with open(self.outputfile, "r") as file:
+                input_data = file.read()
+                outputdict = json.loads(input_data)
+        else:
+            outputdict = []
+
+        outputimagedict = {
+            "image": self.filename,
+            "annotations": []
+        }
+
+        for shape in self.shapes:
+            points = shape["points"]
+
+            x1 = points[0][0]
+            y1 = points[0][1]
+            x2 = points[1][0]
+            y2 = points[2][1]
+
+            height, width, x, y = self.calculate_coordinates(x1, x2, y1, y2)
+
+            shapedict = {
+                "label": shape["label"],
+                "coordinates": {
+                    "x": x,
+                    "y": y,
+                    "width": width,
+                    "height": height
+                }
+            }
+            outputimagedict["annotations"].append(shapedict)
+
+        # check if image already in output
+        exists = False
+        for i in range(0, len(outputdict)):
+            if outputdict[i]["image"] == outputimagedict["image"]:
+                exists = True
+                outputdict[i] = outputimagedict
+                break
+
+        if not exists:
+            outputdict.append(outputimagedict)
+
+        Path(self.outputfile).write_text(json.dumps(outputdict), ENCODE_METHOD)
+
+    def calculate_coordinates(self, x1, x2, y1, y2):
+        if x1 < x2:
+            xmin = x1
+            xmax = x2
+        else:
+            xmin = x2
+            xmax = x1
+        if y1 < y2:
+            ymin = y1
+            ymax = y2
+        else:
+            ymin = y2
+            ymax = y1
+        width = xmax - xmin
+        if width < 0:
+            width = width * -1
+        height = ymax - ymin
+        # x and y from center of rect
+        x = xmin + width / 2
+        y = ymin + height / 2
+        return height, width, x, y
+
+
+class CreateMLReader:
+    def __init__(self, jsonpath, filepath):
+        self.jsonpath = jsonpath
+        self.shapes = []
+        self.verified = False
+        self.filename = filepath.split("/")[-1:][0]
+        try:
+            self.parse_json()
+        except ValueError:
+            print("JSON decoding failed")
+
+    def parse_json(self):
+        with open(self.jsonpath, "r") as file:
+            inputdata = file.read()
+
+        outputdict = json.loads(inputdata)
+        self.verified = True
+
+        if len(self.shapes) > 0:
+            self.shapes = []
+        for image in outputdict:
+            if image["image"] == self.filename:
+                for shape in image["annotations"]:
+                    self.add_shape(shape["label"], shape["coordinates"])
+
+    def add_shape(self, label, bndbox):
+        xmin = bndbox["x"] - (bndbox["width"] / 2)
+        ymin = bndbox["y"] - (bndbox["height"] / 2)
+
+        xmax = bndbox["x"] + (bndbox["width"] / 2)
+        ymax = bndbox["y"] + (bndbox["height"] / 2)
+
+        points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
+        self.shapes.append((label, points, None, None, True))
+
+    def get_shapes(self):
+        return self.shapes

libs/labelFile.py

@@ -10,6 +10,8 @@ from base64 import b64encode, b64decode
 from libs.pascal_voc_io import PascalVocWriter
 from libs.yolo_io import YOLOWriter
 from libs.pascal_voc_io import XML_EXT
+from libs.create_ml_io import CreateMLWriter
+from libs.create_ml_io import JSON_EXT
 from enum import Enum
 import os.path
 import sys
@@ -18,6 +20,7 @@ import sys
 class LabelFileFormat(Enum):
     PASCAL_VOC= 1
     YOLO = 2
+    CREATE_ML = 3
 
 
 class LabelFileError(Exception):
@@ -35,6 +38,23 @@ class LabelFile(object):
         self.imageData = None
         self.verified = False
 
+    def saveCreateMLFormat(self, filename, shapes, imagePath, imageData, classList, lineColor=None, fillColor=None, databaseSrc=None):
+        imgFolderPath = os.path.dirname(imagePath)
+        imgFolderName = os.path.split(imgFolderPath)[-1]
+        imgFileName = os.path.basename(imagePath)
+        outputFilePath = "/".join(filename.split("/")[:-1])
+        outputFile = outputFilePath + "/" + imgFolderName + JSON_EXT
+
+        image = QImage()
+        image.load(imagePath)
+        imageShape = [image.height(), image.width(),
+                      1 if image.isGrayscale() else 3]
+        writer = CreateMLWriter(imgFolderName, imgFileName,
+                                imageShape, shapes, outputFile, localimgpath=imagePath)
+        writer.verified = self.verified
+        writer.write()
+
+
     def savePascalVocFormat(self, filename, shapes, imagePath, imageData,
                             lineColor=None, fillColor=None, databaseSrc=None):
         imgFolderPath = os.path.dirname(imagePath)

resources.qrc

@@ -20,6 +20,7 @@
 <file alias="save">resources/icons/save.png</file>
 <file alias="format_voc">resources/icons/format_voc.png</file>
 <file alias="format_yolo">resources/icons/format_yolo.png</file>
+<file alias="format_createml">resources/icons/format_createml.png</file>
 <file alias="save-as">resources/icons/save-as.png</file>
 <file alias="color">resources/icons/color.png</file>
 <file alias="color_line">resources/icons/color_line.png</file>

tests/test_io.py

@@ -28,5 +28,72 @@ class TestPascalVocRW(unittest.TestCase):
         self.assertEqual(face[0], 'face')
         self.assertEqual(face[1], [(113, 40), (450, 40), (450, 403), (113, 403)])
 
+
+class TestCreateMLRW(unittest.TestCase):
+
+    def test_a_write(self):
+        dir_name = os.path.abspath(os.path.dirname(__file__))
+        libs_path = os.path.join(dir_name, '..', 'libs')
+        sys.path.insert(0, libs_path)
+        from create_ml_io import CreateMLWriter
+
+        person = {'label': 'person', 'points': ((65, 45), (420, 45), (420, 512), (65, 512))}
+        face = {'label': 'face', 'points': ((245, 250), (350, 250), (350, 365), (245, 365))}
+
+        expected_width = 105    # 350-245 -> create_ml_io.py ll 46
+        expected_height = 115   # 365-250 -> create_ml_io.py ll 49
+        expected_x = 297.5      # 245+105/2 -> create_ml_io.py ll 53
+        expected_y = 307.5      # 250+115/2 > create_ml_io.py ll 54
+
+        shapes = [person, face]
+        output_file = dir_name + "/tests.json"
+
+        writer = CreateMLWriter('tests', 'test.512.512.bmp', (512, 512, 1), shapes, output_file,
+                                localimgpath='tests/test.512.512.bmp')
+        writer.write()
+
+        # check written json
+        with open(output_file, "r") as file:
+            inputdata = file.read()
+
+        import json
+        data_dict = json.loads(inputdata)[0]
+
+        self.assertEqual('test.512.512.bmp', data_dict['image'], 'filename not correct in .json')
+        self.assertEqual(2, len(data_dict['annotations']), 'outputfile contains to less annotations')
+        face = data_dict['annotations'][1]
+        self.assertEqual('face', face['label'], 'labelname is wrong')
+        face_coords = face['coordinates']
+        self.assertEqual(expected_width, face_coords['width'], 'calculated width is wrong')
+        self.assertEqual(expected_height, face_coords['height'], 'calculated height is wrong')
+        self.assertEqual(expected_x, face_coords['x'], 'calculated x is wrong')
+        self.assertEqual(expected_y, face_coords['y'], 'calculated y is wrong')
+
+    def test_b_read(self):
+        dir_name = os.path.abspath(os.path.dirname(__file__))
+        libs_path = os.path.join(dir_name, '..', 'libs')
+        sys.path.insert(0, libs_path)
+        from create_ml_io import CreateMLReader
+
+        output_file = dir_name + "/tests.json"
+        reader = CreateMLReader(output_file, 'tests/test.512.512.bmp')
+        shapes = reader.get_shapes()
+        face = shapes[1]
+
+        self.assertEqual(2, len(shapes), 'shape count is wrong')
+        self.assertEqual('face', face[0], 'label is wrong')
+
+        face_coords = face[1]
+        xmin = face_coords[0][0]
+        xmax = face_coords[1][0]
+        ymin = face_coords[0][1]
+        ymax = face_coords[2][1]
+
+        self.assertEqual(245, xmin, 'xmin is wrong')
+        self.assertEqual(350, xmax, 'xmax is wrong')
+        self.assertEqual(250, ymin, 'ymin is wrong')
+        self.assertEqual(365, ymax, 'ymax is wrong')
+
+
 if __name__ == '__main__':
     unittest.main()