From 3171b82e292edb743414d8688c7e9fed6e51cdba Mon Sep 17 00:00:00 2001 From: Licht Takeuchi Date: Sat, 27 Mar 2021 13:52:54 +0900 Subject: [PATCH] feat: add YOLO export https://github.com/microsoft/VoTT/issues/1034 --- src/common/localization/en-us.ts | 3 + src/common/localization/es-cl.ts | 3 + src/common/localization/ja.ts | 3 + src/common/localization/ko-kr.ts | 3 + src/common/localization/zh-ch.ts | 3 + src/common/localization/zh-tw.ts | 3 + src/common/strings.ts | 3 + src/providers/export/yolo.json | 22 ++ src/providers/export/yolo.test.ts | 335 ++++++++++++++++++++++++++++++ src/providers/export/yolo.ts | 188 +++++++++++++++++ src/providers/export/yolo.ui.json | 2 + src/registerProviders.ts | 6 + 12 files changed, 574 insertions(+) create mode 100644 src/providers/export/yolo.json create mode 100644 src/providers/export/yolo.test.ts create mode 100644 src/providers/export/yolo.ts create mode 100644 src/providers/export/yolo.ui.json diff --git a/src/common/localization/en-us.ts b/src/common/localization/en-us.ts index 82f0438b0e..c4f25a0564 100644 --- a/src/common/localization/en-us.ts +++ b/src/common/localization/en-us.ts @@ -423,6 +423,9 @@ export const english: IAppStrings = { description: "Whether or not to include unassigned tags in exported data", }, }, + yolo: { + displayName: "YOLO", + }, cntk: { displayName: "Microsoft Cognitive Toolkit (CNTK)", }, diff --git a/src/common/localization/es-cl.ts b/src/common/localization/es-cl.ts index 54f782187d..e13e77bf0b 100644 --- a/src/common/localization/es-cl.ts +++ b/src/common/localization/es-cl.ts @@ -426,6 +426,9 @@ export const spanish: IAppStrings = { description: "Si se incluyen o no etiquetas no asignadas en los datos exportados", }, }, + yolo: { + displayName: "YOLO", + }, cntk: { displayName: "Microsoft Cognitive Toolkit (CNTK)", }, diff --git a/src/common/localization/ja.ts b/src/common/localization/ja.ts index a4698279f6..6aeee2e519 100644 --- a/src/common/localization/ja.ts +++ b/src/common/localization/ja.ts @@ -435,6 +435,9 @@ export const japanese: IAppStrings = { // Whether or not to include unassigned tags in exported data" }, }, + yolo: { + displayName: "YOLO", // YOLO + }, cntk: { displayName: "Microsoft Cognitive Toolkit（CNTK）", // Microsoft Cognitive Toolkit (CNTK)" }, diff --git a/src/common/localization/ko-kr.ts b/src/common/localization/ko-kr.ts index 8eb0992552..ebedc339ba 100644 --- a/src/common/localization/ko-kr.ts +++ b/src/common/localization/ko-kr.ts @@ -436,6 +436,9 @@ export const korean: IAppStrings = { // Whether or not to include unassigned tags in exported data" }, }, + yolo: { + displayName: "YOLO", // YOLO + }, cntk: { displayName: "Microsoft Cognitive Toolkit（CNTK）", // Microsoft Cognitive Toolkit (CNTK)" }, diff --git a/src/common/localization/zh-ch.ts b/src/common/localization/zh-ch.ts index 28c3edb4cf..83cfa41504 100644 --- a/src/common/localization/zh-ch.ts +++ b/src/common/localization/zh-ch.ts @@ -434,6 +434,9 @@ export const chinese: IAppStrings = { description: "是否在导出的数据中包括未被分配的标签", // Whether or not to include unassigned tags in exported data }, }, + yolo: { + displayName: "YOLO", // YOLO + }, cntk: { displayName: "Microsoft Cognitive Toolkit（CNTK)", // Microsoft Cognitive Toolkit (CNTK) }, diff --git a/src/common/localization/zh-tw.ts b/src/common/localization/zh-tw.ts index 18ed24bf17..6871a8b983 100644 --- a/src/common/localization/zh-tw.ts +++ b/src/common/localization/zh-tw.ts @@ -439,6 +439,9 @@ export const chinesetw: IAppStrings = { description: "是否在已匯出的數據中包括未指定的標記", // Whether or not to include unassigned tags in exported data }, }, + yolo: { + displayName: "YOLO", // YOLO + }, cntk: { displayName: "Microsoft Cognitive Toolkit（CNTK)", // Microsoft Cognitive Toolkit (CNTK) }, diff --git a/src/common/strings.ts b/src/common/strings.ts index 4234478c39..83445d3bd4 100644 --- a/src/common/strings.ts +++ b/src/common/strings.ts @@ -424,6 +424,9 @@ export interface IAppStrings { description: string, }, }, + yolo: { + displayName: string, + }, cntk: { displayName: string, }, diff --git a/src/providers/export/yolo.json b/src/providers/export/yolo.json new file mode 100644 index 0000000000..50cc7864fe --- /dev/null +++ b/src/providers/export/yolo.json @@ -0,0 +1,22 @@ +{ + "type": "object", + "title": "${strings.export.providers.yolo.displayName}", + "properties": { + "assetState": { + "type": "string", + "title": "${strings.export.providers.common.properties.assetState.title}", + "description": "${strings.export.providers.common.properties.assetState.description}", + "enum": [ + "all", + "visited", + "tagged" + ], + "default": "visited", + "enumNames": [ + "${strings.export.providers.common.properties.assetState.options.all}", + "${strings.export.providers.common.properties.assetState.options.visited}", + "${strings.export.providers.common.properties.assetState.options.tagged}" + ] + } + } +} diff --git a/src/providers/export/yolo.test.ts b/src/providers/export/yolo.test.ts new file mode 100644 index 0000000000..790b442772 --- /dev/null +++ b/src/providers/export/yolo.test.ts @@ -0,0 +1,335 @@ +import { YOLOExportProvider} from "./yolo"; +import { ExportAssetState } from "./exportProvider"; +import registerProviders from "../../registerProviders"; +import { ExportProviderFactory } from "./exportProviderFactory"; +import { + IAssetMetadata, AssetState, IExportProviderOptions, +} from "../../models/applicationState"; +import MockFactory from "../../common/mockFactory"; + +jest.mock("../../services/assetService"); +import { AssetService } from "../../services/assetService"; + +jest.mock("../storage/localFileSystemProxy"); +import { LocalFileSystemProxy } from "../storage/localFileSystemProxy"; +import registerMixins from "../../registerMixins"; +import HtmlFileReader from "../../common/htmlFileReader"; +import { appInfo } from "../../common/appInfo"; +import { AssetProviderFactory } from "../storage/assetProviderFactory"; + +registerMixins(); + +describe("YOLO Format Export Provider", () => { + const testAssets = MockFactory.createTestAssets(10, 1); + const baseTestProject = MockFactory.createTestProject("Test Project"); + baseTestProject.assets = { + "asset-1": MockFactory.createTestAsset("1", AssetState.Tagged), + "asset-2": MockFactory.createTestAsset("2", AssetState.Tagged), + "asset-3": MockFactory.createTestAsset("3", AssetState.Visited), + }; + baseTestProject.sourceConnection = MockFactory.createTestConnection("test", "localFileSystemProxy"); + baseTestProject.targetConnection = MockFactory.createTestConnection("test", "localFileSystemProxy"); + + HtmlFileReader.getAssetArray = jest.fn(() => { + return Promise.resolve(new Uint8Array([1, 2, 3]).buffer); + }); + + beforeAll(() => { + AssetProviderFactory.create = jest.fn(() => { + return { + getAssets: jest.fn(() => Promise.resolve(testAssets)), + }; + }); + }); + + beforeEach(() => { + registerProviders(); + }); + + it("Is defined", () => { + expect(YOLOExportProvider).toBeDefined(); + }); + + it("Can be instantiated through the factory", () => { + const options: IExportProviderOptions = { + assetState: ExportAssetState.All, + }; + const exportProvider = ExportProviderFactory.create("yolo", baseTestProject, options); + expect(exportProvider).not.toBeNull(); + expect(exportProvider).toBeInstanceOf(YOLOExportProvider); + }); + + describe("Export variations", () => { + beforeEach(() => { + const assetServiceMock = AssetService as jest.Mocked; + assetServiceMock.prototype.getAssetMetadata = jest.fn((asset) => { + const isOdd = Number(asset.id.split("-")[1]) % 2 === 0; + const mockTag1 = MockFactory.createTestTag(isOdd ? "1" : "2"); + const mockTag2 = MockFactory.createTestTag(isOdd ? "2" : "1"); + const mockTag3 = MockFactory.createTestTag("3"); + const mockRegion1 = MockFactory.createTestRegion("region-1", [mockTag1.name]); + const mockRegion2 = MockFactory.createTestRegion("region-2", [mockTag2.name]); + const mockRegion3 = MockFactory.createTestRegion("region-3", [mockTag3.name]); + + const assetMetadata: IAssetMetadata = { + asset, + regions: [mockRegion1, mockRegion2, mockRegion3], + version: appInfo.version, + }; + + return Promise.resolve(assetMetadata); + }); + + const storageProviderMock = LocalFileSystemProxy as jest.Mock; + storageProviderMock.mockClear(); + }); + + it("Exports all assets", async () => { + const options: IExportProviderOptions = { + assetState: ExportAssetState.All, + }; + + const testProject = { ...baseTestProject }; + testProject.tags = MockFactory.createTestTags(4); + + const exportProvider = new YOLOExportProvider(testProject, options); + await exportProvider.export(); + + const storageProviderMock = LocalFileSystemProxy as any; + const createContainerCalls = storageProviderMock.mock.instances[0].createContainer.mock.calls; + + expect(createContainerCalls.length).toEqual(3); + expect(createContainerCalls[1][0].endsWith("/JPEGImages")).toEqual(true); + expect(createContainerCalls[2][0].endsWith("/labels")).toEqual(true); + + const writeBinaryCalls = storageProviderMock.mock.instances[0].writeBinary.mock.calls; + expect(writeBinaryCalls.length).toEqual(testAssets.length); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < testAssets.length; i++) { + expect(writeBinaryCalls[i][0].endsWith(`/JPEGImages/Asset ${i + 1}.jpg`)).toEqual(true); + } + + const writeTextFileCalls = storageProviderMock.mock.instances[0].writeText.mock.calls as any[]; + // We write an annotation txt file per asset and 1 names file + expect(writeTextFileCalls.length).toEqual(testAssets.length + 1); + + const assetServiceMock = AssetService as any; + const assetServiceLastIndex = assetServiceMock.mock.instances.length - 1; + const getAssetMetadataResults = + await Promise.all( + (assetServiceMock.mock.instances[assetServiceLastIndex] + .getAssetMetadata.mock.results as any[]).map((x) => x.value), + ); + expect(getAssetMetadataResults.length).toEqual(testAssets.length); + + // tslint:disable-next-line:prefer-for-of + for (let assetIndex = 0; assetIndex < testAssets.length; assetIndex++) { + const asset = testAssets[assetIndex]; + const assetId = assetIndex + 1; + const labelIndex = + writeTextFileCalls.findIndex((args) => args[0].endsWith(`/labels/Asset ${assetId}.txt`)); + expect(labelIndex).toBeGreaterThanOrEqual(0); + const labelRecords = writeTextFileCalls[labelIndex][1].split("\n"); + + const isAssetIdOdd = assetId % 2 === 0; + + const assetMetadataResultIndex = + getAssetMetadataResults.findIndex((args) => args.asset.id === `asset-${assetId}`); + expect(assetMetadataResultIndex).toBeGreaterThanOrEqual(0); + const regions = getAssetMetadataResults[assetMetadataResultIndex].regions; + expect(labelRecords.length).toEqual(regions.length); + + // tslint:disable-next-line:prefer-for-of + for (let recordIndex = 0; recordIndex < labelRecords.length; recordIndex++) { + const fields = labelRecords[recordIndex].split(" "); + expect(fields.length).toEqual(5); + + const bbox = regions[recordIndex].boundingBox; + + const oddId = recordIndex === 0 ? 1 : 2; + const evenId = recordIndex === 0 ? 2 : 1; + + expect(fields[0]).toEqual((recordIndex === 2 ? 3 : isAssetIdOdd ? oddId : evenId).toString()); + expect(fields[1]).toEqual(((2 * bbox.left + bbox.width) / (2 * asset.size.width)).toString()); + expect(fields[2]).toEqual(((2 * bbox.top + bbox.height) / (2 * asset.size.height)).toString()); + expect(fields[3]).toEqual((bbox.width / asset.size.width).toString()); + expect(fields[4]).toEqual((bbox.height / asset.size.height).toString()); + } + } + + expect(writeTextFileCalls[0][0].endsWith("yolo.names")).toEqual(true); + const nameRecords = writeTextFileCalls[0][1].split("\n"); + expect(nameRecords.length).toEqual(testProject.tags.length); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < testProject.tags.length; i++) { + expect(nameRecords[i]).toEqual(`Tag ${i}`); + } + }); + + it("Exports only visited assets (includes tagged)", async () => { + const options: IExportProviderOptions = { + assetState: ExportAssetState.Visited, + }; + + const testProject = { ...baseTestProject }; + testProject.tags = MockFactory.createTestTags(4); + + const exportProvider = new YOLOExportProvider(testProject, options); + await exportProvider.export(); + + const storageProviderMock = LocalFileSystemProxy as any; + const createContainerCalls = storageProviderMock.mock.instances[0].createContainer.mock.calls; + + expect(createContainerCalls.length).toEqual(3); + expect(createContainerCalls[1][0].endsWith("/JPEGImages")).toEqual(true); + expect(createContainerCalls[2][0].endsWith("/labels")).toEqual(true); + + const numWriteBinaryCalls = Object.keys(testProject.assets).length; + + const writeBinaryCalls = storageProviderMock.mock.instances[0].writeBinary.mock.calls; + expect(writeBinaryCalls.length).toEqual(numWriteBinaryCalls); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < numWriteBinaryCalls; i++) { + expect(writeBinaryCalls[i][0].endsWith(`/JPEGImages/Asset ${i + 1}.jpg`)).toEqual(true); + } + + const writeTextFileCalls = storageProviderMock.mock.instances[0].writeText.mock.calls as any[]; + expect(writeTextFileCalls.length).toEqual(numWriteBinaryCalls + 1); + + const assetServiceMock = AssetService as any; + const assetServiceLastIndex = assetServiceMock.mock.instances.length - 1; + const getAssetMetadataResults = + await Promise.all( + (assetServiceMock.mock.instances[assetServiceLastIndex] + .getAssetMetadata.mock.results as any[]).map((x) => x.value), + ); + expect(getAssetMetadataResults.length).toEqual(numWriteBinaryCalls); + + // tslint:disable-next-line:prefer-for-of + for (let assetIndex = 0; assetIndex < numWriteBinaryCalls; assetIndex++) { + const asset = testAssets[assetIndex]; + const assetId = assetIndex + 1; + const labelIndex = + writeTextFileCalls.findIndex((args) => args[0].endsWith(`/labels/Asset ${assetId}.txt`)); + expect(labelIndex).toBeGreaterThanOrEqual(0); + const labelRecords = writeTextFileCalls[labelIndex][1].split("\n"); + + const isAssetIdOdd = assetId % 2 === 0; + + const assetMetadataResultIndex = + getAssetMetadataResults.findIndex((args) => args.asset.id === `asset-${assetId}`); + expect(assetMetadataResultIndex).toBeGreaterThanOrEqual(0); + const regions = getAssetMetadataResults[assetMetadataResultIndex].regions; + expect(labelRecords.length).toEqual(regions.length); + + // tslint:disable-next-line:prefer-for-of + for (let recordIndex = 0; recordIndex < labelRecords.length; recordIndex++) { + const fields = labelRecords[recordIndex].split(" "); + expect(fields.length).toEqual(5); + + const bbox = regions[recordIndex].boundingBox; + + const oddId = recordIndex === 0 ? 1 : 2; + const evenId = recordIndex === 0 ? 2 : 1; + + expect(fields[0]).toEqual((recordIndex === 2 ? 3 : isAssetIdOdd ? oddId : evenId).toString()); + expect(fields[1]).toEqual(((2 * bbox.left + bbox.width) / (2 * asset.size.width)).toString()); + expect(fields[2]).toEqual(((2 * bbox.top + bbox.height) / (2 * asset.size.height)).toString()); + expect(fields[3]).toEqual((bbox.width / asset.size.width).toString()); + expect(fields[4]).toEqual((bbox.height / asset.size.height).toString()); + } + } + + expect(writeTextFileCalls[0][0].endsWith("yolo.names")).toEqual(true); + const nameRecords = writeTextFileCalls[0][1].split("\n"); + expect(nameRecords.length).toEqual(testProject.tags.length); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < testProject.tags.length; i++) { + expect(nameRecords[i]).toEqual(`Tag ${i}`); + } + }); + + it("Exports only tagged assets", async () => { + const options: IExportProviderOptions = { + assetState: ExportAssetState.Tagged, + }; + + const testProject = { ...baseTestProject }; + testProject.tags = MockFactory.createTestTags(4); + + const exportProvider = new YOLOExportProvider(testProject, options); + await exportProvider.export(); + + const storageProviderMock = LocalFileSystemProxy as any; + const createContainerCalls = storageProviderMock.mock.instances[0].createContainer.mock.calls; + + expect(createContainerCalls.length).toEqual(3); + expect(createContainerCalls[1][0].endsWith("/JPEGImages")).toEqual(true); + expect(createContainerCalls[2][0].endsWith("/labels")).toEqual(true); + + const numWriteBinaryCalls = 2; + + const writeBinaryCalls = storageProviderMock.mock.instances[0].writeBinary.mock.calls; + expect(writeBinaryCalls.length).toEqual(numWriteBinaryCalls); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < 2; i++) { + expect(writeBinaryCalls[i][0].endsWith(`/JPEGImages/Asset ${i + 1}.jpg`)).toEqual(true); + } + + const writeTextFileCalls = storageProviderMock.mock.instances[0].writeText.mock.calls as any[]; + expect(writeTextFileCalls.length).toEqual(numWriteBinaryCalls + 1); + + const assetServiceMock = AssetService as any; + const assetServiceLastIndex = assetServiceMock.mock.instances.length - 1; + const getAssetMetadataResults = + await Promise.all( + (assetServiceMock.mock.instances[assetServiceLastIndex] + .getAssetMetadata.mock.results as any[]).map((x) => x.value), + ); + expect(getAssetMetadataResults.length).toEqual(numWriteBinaryCalls); + + // tslint:disable-next-line:prefer-for-of + for (let assetIndex = 0; assetIndex < numWriteBinaryCalls; assetIndex++) { + const asset = testAssets[assetIndex]; + const assetId = assetIndex + 1; + const labelIndex = + writeTextFileCalls.findIndex((args) => args[0].endsWith(`/labels/Asset ${assetId}.txt`)); + expect(labelIndex).toBeGreaterThanOrEqual(0); + const labelRecords = writeTextFileCalls[labelIndex][1].split("\n"); + + const isAssetIdOdd = assetId % 2 === 0; + + const assetMetadataResultIndex = + getAssetMetadataResults.findIndex((args) => args.asset.id === `asset-${assetId}`); + expect(assetMetadataResultIndex).toBeGreaterThanOrEqual(0); + const regions = getAssetMetadataResults[assetMetadataResultIndex].regions; + expect(labelRecords.length).toEqual(regions.length); + + // tslint:disable-next-line:prefer-for-of + for (let recordIndex = 0; recordIndex < labelRecords.length; recordIndex++) { + const fields = labelRecords[recordIndex].split(" "); + expect(fields.length).toEqual(5); + + const bbox = regions[recordIndex].boundingBox; + + const oddId = recordIndex === 0 ? 1 : 2; + const evenId = recordIndex === 0 ? 2 : 1; + + expect(fields[0]).toEqual((recordIndex === 2 ? 3 : isAssetIdOdd ? oddId : evenId).toString()); + expect(fields[1]).toEqual(((2 * bbox.left + bbox.width) / (2 * asset.size.width)).toString()); + expect(fields[2]).toEqual(((2 * bbox.top + bbox.height) / (2 * asset.size.height)).toString()); + expect(fields[3]).toEqual((bbox.width / asset.size.width).toString()); + expect(fields[4]).toEqual((bbox.height / asset.size.height).toString()); + } + } + + expect(writeTextFileCalls[0][0].endsWith("yolo.names")).toEqual(true); + const nameRecords = writeTextFileCalls[0][1].split("\n"); + expect(nameRecords.length).toEqual(testProject.tags.length); + // tslint:disable-next-line:prefer-for-of + for (let i = 0; i < testProject.tags.length; i++) { + expect(nameRecords[i]).toEqual(`Tag ${i}`); + } + }); + }); +}); diff --git a/src/providers/export/yolo.ts b/src/providers/export/yolo.ts new file mode 100644 index 0000000000..b9c3bc99a6 --- /dev/null +++ b/src/providers/export/yolo.ts @@ -0,0 +1,188 @@ +import _ from "lodash"; +import { ExportProvider } from "./exportProvider"; +import { IProject, IAssetMetadata, ITag, IExportProviderOptions } from "../../models/applicationState"; +import Guard from "../../common/guard"; +import HtmlFileReader from "../../common/htmlFileReader"; +import json2csv, { Parser } from "json2csv"; +import {ISize} from "vott-react"; + +interface IObjectInfo { + id: number; + xcenter: number; + ycenter: number; + width: number; + height: number; +} + +/** + * @name - YOLO Export Provider + * @description - Exports a project into a YOLO + */ +export class YOLOExportProvider extends ExportProvider { + private labelsInfo = new Map(); + private nameMap = new Map(); + + constructor(project: IProject, options: IExportProviderOptions) { + super(project, options); + Guard.null(options); + } + + /** + * Export project to YOLO + */ + public async export(): Promise { + const allAssets = await this.getAssetsForExport(); + const exportObject: any = { ...this.project }; + exportObject.assets = _.keyBy(allAssets, (assetMetadata) => assetMetadata.asset.id); + + // Create Export Folder + const exportFolderName = `${this.project.name.replace(/\s/g, "-")}-YOLO-export`; + await this.storageProvider.createContainer(exportFolderName); + + await this.exportNames(exportFolderName); + await this.exportImages(exportFolderName, allAssets); + await this.exportLabels(exportFolderName, allAssets); + } + + private async exportImages(exportFolderName: string, allAssets: IAssetMetadata[]) { + // Create JPEGImages Sub Folder + const jpegImagesFolderName = `${exportFolderName}/JPEGImages`; + await this.storageProvider.createContainer(jpegImagesFolderName); + + await allAssets.mapAsync(async (assetMetadata) => { + await this.exportSingleImage(jpegImagesFolderName, assetMetadata); + }); + } + + private async exportSingleImage(jpegImagesFolderName: string, assetMetadata: IAssetMetadata): Promise { + try { + const arrayBuffer = await HtmlFileReader.getAssetArray(assetMetadata.asset); + const buffer = Buffer.from(arrayBuffer); + const imageFileName = assetMetadata.asset.name; + const imageFilePath = `${jpegImagesFolderName}/${imageFileName}`; + + // Write Binary + await this.storageProvider.writeBinary(imageFilePath, buffer); + + const imageSize = await this.getImageSize(arrayBuffer, imageFilePath, assetMetadata); + + // Get Array of all Box shaped tag for the Asset + const tagObjects = this.getAssetTagArray(assetMetadata, imageSize); + const labelFileName = `${imageFileName.substr(0, imageFileName.lastIndexOf(".")) + || imageFileName}.txt`; + this.labelsInfo.set(labelFileName, tagObjects); + } catch (err) { + // Ignore the error at the moment + // TODO: Refactor ExportProvider abstract class export() method + // to return Promise with an object containing + // the number of files successfully exported out of total + console.log(`Error downloading asset ${assetMetadata.asset.path} - ${err}`); + } + } + + private getAssetTagArray(element: IAssetMetadata, imageSize: ISize): IObjectInfo[] { + const tagObjects = []; + element.regions.forEach((region) => { + region.tags.forEach((tagName) => { + const objectInfo: IObjectInfo = { + id: this.nameMap.get(tagName), + xcenter: (2 * region.boundingBox.left + region.boundingBox.width) / (2 * imageSize.width), + ycenter: (2 * region.boundingBox.top + region.boundingBox.height) / (2 * imageSize.height), + width: region.boundingBox.width / imageSize.width, + height: region.boundingBox.height / imageSize.height, + }; + + tagObjects.push(objectInfo); + }); + }); + + return tagObjects; + } + + private async getImageSize(imageBuffer: ArrayBuffer, imageFilePath: string, assetMetadata: IAssetMetadata) + : Promise { + if (assetMetadata.asset.size && + assetMetadata.asset.size.width !== 0 && + assetMetadata.asset.size.height !== 0) { + return assetMetadata.asset.size; + } + + // Get Base64 + const image64 = btoa(new Uint8Array(imageBuffer). + reduce((data, byte) => data + String.fromCharCode(byte), "")); + + if (image64.length < 10) { + // Ignore the error at the moment + // TODO: Refactor ExportProvider abstract class export() method + // to return Promise with an object containing + // the number of files successfully exported out of total + console.log(`Image not valid ${imageFilePath}`); + } else { + const assetProps = await HtmlFileReader.readAssetAttributesWithBuffer(image64); + if (assetProps) { + return assetProps; + } + + console.log(`imageInfo for element ${assetMetadata.asset.name} not found (${assetProps})`); + } + + return {width: 1, height: 1}; + } + + private async exportNames(exportFolderName: string) { + if (!this.project.tags) { + return; + } + + const namesFileName = `${exportFolderName}/yolo.names`; + + let id = 0; + const dataItems = this.project.tags.map((element) => { + const objectName = element.name; + this.nameMap.set(objectName, id++); + return {name: objectName}; + }); + + // Configure CSV options + const csvOptions: json2csv.Options<{}> = { + fields: ["name"], + header: false, + delimiter: "", + quote: "", + doubleQuote: "", + eol: "\n", + }; + const csvParser = new Parser(csvOptions); + const csvData = csvParser.parse(dataItems); + + await this.storageProvider.writeText(namesFileName, csvData); + } + + private async exportLabels(exportFolderName: string, allAssets: IAssetMetadata[]) { + // Create Labels Sub Folder + const labelsFolderName = `${exportFolderName}/labels`; + await this.storageProvider.createContainer(labelsFolderName); + + try { + // Save Labels + await this.labelsInfo.forEachAsync(async (objectsInfo, labelFileName) => { + const labelFilePath = `${labelsFolderName}/${labelFileName}`; + + const csvOptions: json2csv.Options<{}> = { + fields: ["id", "xcenter", "ycenter", "width", "height"], + header: false, + delimiter: " ", + quote: "", + doubleQuote: "", + eol: "\n", + }; + const csvParser = new Parser(csvOptions); + const csvData = csvParser.parse(objectsInfo); + + await this.storageProvider.writeText(labelFilePath, csvData); + }); + } catch (err) { + console.log("Error writing YOLO annotation file"); + } + } +} diff --git a/src/providers/export/yolo.ui.json b/src/providers/export/yolo.ui.json new file mode 100644 index 0000000000..2c63c08510 --- /dev/null +++ b/src/providers/export/yolo.ui.json @@ -0,0 +1,2 @@ +{ +} diff --git a/src/registerProviders.ts b/src/registerProviders.ts index e01209b0e0..03cc5ee207 100644 --- a/src/registerProviders.ts +++ b/src/registerProviders.ts @@ -1,5 +1,6 @@ import { ExportProviderFactory } from "./providers/export/exportProviderFactory"; import { PascalVOCExportProvider } from "./providers/export/pascalVOC"; +import { YOLOExportProvider } from "./providers/export/yolo"; import { TFRecordsExportProvider } from "./providers/export/tensorFlowRecords"; import { VottJsonExportProvider } from "./providers/export/vottJson"; import { CsvExportProvider } from "./providers/export/csv"; @@ -60,6 +61,11 @@ export default function registerProviders() { displayName: strings.export.providers.pascalVoc.displayName, factory: (project, options) => new PascalVOCExportProvider(project, options), }); + ExportProviderFactory.register({ + name: "yolo", + displayName: strings.export.providers.yolo.displayName, + factory: (project, options) => new YOLOExportProvider(project, options), + }); ExportProviderFactory.register({ name: "tensorFlowRecords", displayName: strings.export.providers.tfRecords.displayName,