From 31463497c89928e7d2ea47d3e8bf5a953fdf8aeb Mon Sep 17 00:00:00 2001 From: enicck <37335354+enicck@users.noreply.github.com> Date: Fri, 2 Oct 2020 15:56:07 +0200 Subject: [PATCH] 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 | 54 +++++++++++- libs/constants.py | 1 + libs/create_ml_io.py | 131 ++++++++++++++++++++++++++++ libs/labelFile.py | 20 +++++ resources.qrc | 1 + resources/icons/format_createml.png | Bin 0 -> 4156 bytes tests/test_io.py | 67 ++++++++++++++ 7 files changed, 270 insertions(+), 4 deletions(-) create mode 100644 libs/create_ml_io.py create mode 100644 resources/icons/format_createml.png 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 0000000000000000000000000000000000000000..c08af9426e7888aa05f57c58c6f680442b80339b GIT binary patch literal 4156 zcmb_f2~<-_){ep`AR+P(TD_0KptC1@>B234?N26qp~*73s>S0wElSSOMT2 z>*m9Z4P}y8Fk2g_bu<|?5Dtjw(CF|mu8j zFYt{53lWL8OD_r%>0u2x{1suMJ z!{b7g7U>LLq=*6onSKc&od4Y{SNJteV8Rg5bUp%QiByI(3uH0B_#NOL+Kq7edj-U0@hr6Y{5Mp-&7S zKzIQ{UZj8tY}^IXT&0YLPo@e0x`-$6;qk)070TlqWhfeDiH5Fqr87BPWe{e+n*z}2 zB7g!@)(ws%!O>VB6poBUlW|B3@Qy@&fx7Zo9Cpkfpu|}yfs7;m0~8bui!P%7M=*;? zX7dE$bTDL2I6W9Z@VUV-==VgDsk|_r0CWtpL;b$rkxKOt@YtL%aDc`WLG9@rE)<1Z zJ;7Fsuz|2^3I4m-Oz+mH1bOszxB(dOF6p0BZ5?M?*flVUPNpuv0 z#YTUdPvbEomGxK7|FRTVJSJ%4|CtdmBs3OB#KDmmJOhqqV%cyaiGYO@*?1-v%O(CV^5CvSB4hF2B>Ufy|C$mF+_@0Zw~cM=0$sYx`kc1*4I1|6LVw_W8=pS_Owd%aDW?5UXikl5DOl{l90 zA>l>T!#f|Oaj$>cFwrEQSQt@~qWpheB=7Kuv~C|rUp7_fk$o{69Z$YV8)Wplm$?Pe z&-B%q(B6xWf9RDj@IN%>mxyYS5*TcS9aHwdgy>}qHYN{t7Img%C8NK ziW!wERJWb%%R^k`8i_)wIGqO_E698#FMWEtn&zVuI~vBG1VAuugDlHD4q_j&(j*Gok= z#scn$8V~W7RxYnNXmq-mzOBV^zeIdM_5BTZ4yJHqW&50j&3;nW^P0w(4o~fyKNN|Z zGEdZw4BkF}w_DmBlV&5U-CY!a`>j^d=)EIeX{kpab!UzA=f94&_zcSfip|L>>o)bz z!FG)J8u&^rH?G&u-aN2y=(K<3wXSF7X~T$vB_}3p?#I~luM5uqO?RnZ=7)n*^@9|j zBb{k0TV>Lx8ziX#a=G07;$1|u)PHB-;{n{&J!;OTk29!#HPQC&4c*c)mzN(NI~W)R z>`PNQNzfd0(`a$Loc=@bdwTrkvBzb%SBMXr<)7BN|1^6_7MB|>Y#gN4AX{j$_*|y zgyA;D`NEW>O|&;%9hm}m-;cGFo?>MC(WADW)y6D{_eS*^y_k`+N%gAF3@zU3H~I@2 z=+;~G-z2+B`%mg)Gwtv0_*5Tvy5w49V%^^1(f;jsa2kOV?s_)Y^M)tf8)Z>kS=fr$ zUWX$^fuX||?b+`=PAPuYF^8%%nUFk&>E&)wp9$)t<_?80%*(huRS_Hd{$-*dah{M^ z-IaV|$CJPc_VbzhP0h!7+M=&hR7 z1GRA#KOwA}hEh7ZhW%|HxEQn_5#Y7G_6?}c!J99eYn;=!^j%vWpZ^?740&#H9F?%E zVf(?HB3|(+8%b8R<^h_6>bm?jD)j5lr;@J8tM$wHNtXgeq9MpLj(lHMqn%Z|!?<2Y zccj(3h`TS}Q!y|4(>lzyo^?-gN;jAr=3Uibrp4dVX;nF0DxGw)t~_|`V&q(DO-f?( zbX%(tsVDK{cx#BW-s{2x84i#iAuh5`?CPN82Hy3_E4$y~!!8CE$}lX?u0~_eqjC>> z$kKvH+s-_^Wf1m=3xAK6q0@waULEGlLTCZKBZ<%pEY1CN4(KnTZ1<+;b?)_m?an;m zysKhOlBHXzwqJA$X_0q+`mfOWS^BS6=g$j`KtckAzVgs}MdgNF@91{bPNYelV?F$j zcRxg~482(&9CI37mv2uo8qmfxP=SPN_To`#A@OgQOr$d=N!#2|0eI}ggegbs^|q3q zDk?58vihRVEht)aE4NayNqZk;ZoCU`gRCjW{7wIUY8spL-b6JhJ%hbqeRXx7Eb2i{ z>!i%>692S+-xC9?nDeseL*`>_Ew`YIW%25lSB?~_oSolqY%uZU+@gIM8l?*!LrS<^ z(;?dxw)ahsKAda_awPm3z*zGFlAab_p@J>VPB=y+#+n}*bzE=ZJ`S65BK)@ACbfHe zsI4?2q;`n29{Ot?&Kl^D_R_8y7HX~1cvX-YyWTmwMW3i<)6LFi+Kokc=e}pK~Hwtz*^I+^tmelcO9++jTIy`a5fe2>Pk~~@RN&bZinvA4u|X8us;g{=p)J$5O-?;SU~(XMO=O8Gy|;mwW! literal 0 HcmV?d00001 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()