From 955230d13359d640ace8025f63cc24760117f7f3 Mon Sep 17 00:00:00 2001 From: Jamie Date: Tue, 12 Mar 2024 09:33:02 +0800 Subject: [PATCH] Support Rect exporting to COCO JSON --- README.md | 2 +- src/data/ExportFormatData.ts | 4 + .../__tests__/export/RectLabelExport.test.ts | 20 ++++ src/logic/export/RectLabelsExporter.ts | 91 +++++++++++++++++++ 4 files changed, 116 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index d3cb5d21..6b1b5108 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,7 @@ docker logs make-sense |:-------------:|:---:|:----:|:-------:|:--------:|:---------:|:----------:| | **Point** | ✓ | ✗ | ☐ | ☐ | ☐ | ✗ | | **Line** | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | -| **Rect** | ✓ | ✓ | ✓ | ☐ | ☐ | ✗ | +| **Rect** | ✓ | ✓ | ✓ | ☐ | ✓ | ✗ | | **Polygon** | ☐ | ✗ | ☐ | ✓ | ✓ | ☐ | | **Label** | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | diff --git a/src/data/ExportFormatData.ts b/src/data/ExportFormatData.ts index 540536a5..b6ad076b 100644 --- a/src/data/ExportFormatData.ts +++ b/src/data/ExportFormatData.ts @@ -17,6 +17,10 @@ export const ExportFormatData: ExportFormatDataMap = { { type: AnnotationFormatType.CSV, label: 'Single CSV file.' + }, + { + type: AnnotationFormatType.COCO, + label: 'Single file in COCO JSON format.' } ], [LabelType.POINT]: [ diff --git a/src/logic/__tests__/export/RectLabelExport.test.ts b/src/logic/__tests__/export/RectLabelExport.test.ts index 9e19550e..ee82716c 100644 --- a/src/logic/__tests__/export/RectLabelExport.test.ts +++ b/src/logic/__tests__/export/RectLabelExport.test.ts @@ -181,3 +181,23 @@ describe('RectLabelsExporter wrapRectLabelIntoCSV method', () => { expect(parseFloat(resultImageHeight)).toBe(1080) }) }) + +describe('RectLabelsExporter getCOCOArea method', () => { + it('should produce correct area', () => { + const labelRect: LabelRect = { + id: 'label-rect-000', + labelId: 'label-002', + rect: { + x: 444, + y: 998, + width: 90, + height: 82 + }, + isVisible: true, + isCreatedByAI: false, + status: LabelStatus.ACCEPTED, + suggestedLabel: 'label-000' + } + expect(RectLabelsExporter.getCOCOArea(labelRect.rect)).toBe(7380); + }); +}); \ No newline at end of file diff --git a/src/logic/export/RectLabelsExporter.ts b/src/logic/export/RectLabelsExporter.ts index 151d4513..5a1f30e8 100644 --- a/src/logic/export/RectLabelsExporter.ts +++ b/src/logic/export/RectLabelsExporter.ts @@ -12,6 +12,17 @@ import {ISize} from '../../interfaces/ISize'; import {NumberUtil} from '../../utils/NumberUtil'; import {RectUtil} from '../../utils/RectUtil'; import {Settings} from '../../settings/Settings'; +import {flatten} from "lodash"; +import {IRect} from '../../interfaces/IRect'; +import {COCOExporter} from './polygon/COCOExporter'; +import { + COCOAnnotation, COCOBBox, + COCOCategory, + COCOImage, + COCOInfo, + COCOObject, + COCOSegmentation +} from '../../data/labels/COCO'; export class RectLabelsExporter { public static export(exportFormatType: AnnotationFormatType): void { @@ -25,6 +36,9 @@ export class RectLabelsExporter { case AnnotationFormatType.CSV: RectLabelsExporter.exportAsCSV(); break; + case AnnotationFormatType.COCO: + RectLabelsExporter.exportAsCOCO(); + break; default: return; } @@ -226,4 +240,81 @@ export class RectLabelsExporter { labelRect, labelNames, imageSize, imageData.fileData.name)); return labelRectsString.join('\n'); } + + private static exportAsCOCO(): void { + try { + const imagesData: ImageData[] = LabelsSelector.getImagesData(); + const labelNames: LabelName[] = LabelsSelector.getLabelNames(); + const projectName: string = GeneralSelector.getProjectName(); + const COCOObject: COCOObject = RectLabelsExporter.mapImagesDataToCOCOObject(imagesData, labelNames, projectName); + const content: string = JSON.stringify(COCOObject); + const fileName: string = `${ExporterUtil.getExportFileName()}.json`; + ExporterUtil.saveAs(content, fileName); + } catch (error) { + // TODO + throw new Error(error as string); + } + } + + private static mapImagesDataToCOCOObject( + imagesData: ImageData[], + labelNames: LabelName[], + projectName: string + ): COCOObject { + return { + "info": COCOExporter.getInfoComponent(projectName), + "images": RectLabelsExporter.getImagesComponent(imagesData), + "annotations": RectLabelsExporter.getAnnotationsComponent(imagesData, labelNames), + "categories":COCOExporter.getCategoriesComponent(labelNames) + } + } + + public static getImagesComponent(imagesData: ImageData[]): COCOImage[] { + return imagesData + .filter((imagesData: ImageData) => imagesData.loadStatus) + .filter((imagesData: ImageData) => imagesData.labelRects.length !== 0) + .map((imageData: ImageData, index: number) => { + const image: HTMLImageElement = ImageRepository.getById(imageData.id); + return { + "id": index + 1, + "width": image.width, + "height": image.height, + "file_name": imageData.fileData.name + } + }) + } + + public static getAnnotationsComponent(imagesData: ImageData[], labelNames: LabelName[]): COCOAnnotation[] { + const labelsMap: LabelDataMap = COCOExporter.mapLabelsData(labelNames); + let id = 0; + const annotations: COCOAnnotation[][] = imagesData + .filter((imagesData: ImageData) => imagesData.loadStatus) + .filter((imagesData: ImageData) => imagesData.labelRects.length !== 0) + .map((imageData: ImageData, index: number) => { + return imageData.labelRects.map((labelRects: LabelRect) => { + return { + "id": id++, + "iscrowd": 0, + "image_id": index + 1, + "category_id": labelsMap[labelRects.labelId], + "segmentation": RectLabelsExporter.getCOCOSegmentation(labelRects.rect), + "bbox": RectLabelsExporter.getCOCOBbox(labelRects.rect), + "area": RectLabelsExporter.getCOCOArea(labelRects.rect) + } + }) + }) + return flatten(annotations); + } + + public static getCOCOSegmentation(rect: IRect): COCOSegmentation { + return []; + } + + public static getCOCOBbox(rect: IRect): COCOBBox { + return [rect.x, rect.y, rect.width, rect.height]; + } + + public static getCOCOArea(rect: IRect): number { + return rect.height*rect.width; + } }