diff --git a/labelImg.py b/labelImg.py
index e386771a..d04a349c 100755
--- a/labelImg.py
+++ b/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.')
@@ -834,7 +855,12 @@ class MainWindow(QMainWindow, WindowMixin):
if annotationFilePath[-4:].lower() != ".txt":
annotationFilePath += TXT_EXT
self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
- self.lineColor.getRgb(), self.fillColor.getRgb())
+ 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:
diff --git a/libs/constants.py b/libs/constants.py
index 8fa4b275..1efda037 100644
--- a/libs/constants.py
+++ b/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'
diff --git a/libs/create_ml_io.py b/libs/create_ml_io.py
new file mode 100644
index 00000000..0d078146
--- /dev/null
+++ b/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
diff --git a/libs/labelFile.py b/libs/labelFile.py
index b335f818..b366d45e 100644
--- a/libs/labelFile.py
+++ b/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)
diff --git a/resources.qrc b/resources.qrc
index 37d2f2d2..bccd75aa 100644
--- a/resources.qrc
+++ b/resources.qrc
@@ -20,6 +20,7 @@
resources/icons/save.png
resources/icons/format_voc.png
resources/icons/format_yolo.png
+resources/icons/format_createml.png
resources/icons/save-as.png
resources/icons/color.png
resources/icons/color_line.png
diff --git a/resources/icons/format_createml.png b/resources/icons/format_createml.png
new file mode 100644
index 00000000..c08af942
Binary files /dev/null and b/resources/icons/format_createml.png differ
diff --git a/tests/test_io.py b/tests/test_io.py
index 7067b403..7a6d022e 100644
--- a/tests/test_io.py
+++ b/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()