feat(frontend): supports exporting pascal voc and tfrecord format files

opendatalab · Nov 5, 2024 · a81cd8f · a81cd8f
1 parent c1e56fe
commit a81cd8f
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 16 deletions.
diff --git a/apps/frontend/src/api/services/samples.ts b/apps/frontend/src/api/services/samples.ts
@@ -83,7 +83,16 @@ export async function updateSampleAnnotationResult(
 export async function outputSample(taskId: number, sampleIds: number[], activeTxt: ExportType) {
   let res;
 
-  if ([ExportType.MASK, ExportType.LABEL_ME, ExportType.YOLO, ExportType.CSV].includes(activeTxt)) {
+  if (
+    [
+      ExportType.MASK,
+      ExportType.LABEL_ME,
+      ExportType.YOLO,
+      ExportType.CSV,
+      ExportType.TF_RECORD,
+      ExportType.PASCAL_VOC,
+    ].includes(activeTxt)
+  ) {
     res = await request.post(
       `/v1/tasks/${taskId}/samples/export`,
       {
@@ -98,7 +107,7 @@ export async function outputSample(taskId: number, sampleIds: number[], activeTx
       },
     );
   } else {
-    await request.post(
+    res = await request.post(
       `/v1/tasks/${taskId}/samples/export`,
       {
         sample_ids: sampleIds,
@@ -124,10 +133,15 @@ export async function outputSample(taskId: number, sampleIds: number[], activeTx
     case ExportType.COCO:
       filename = filename + '.json';
       break;
+    case ExportType.XML:
+      filename = filename + '.xml';
+      break;
     case ExportType.MASK:
     case ExportType.CSV:
     case ExportType.LABEL_ME:
     case ExportType.YOLO:
+    case ExportType.TF_RECORD:
+    case ExportType.PASCAL_VOC:
       url = window.URL.createObjectURL(data as any);
       break;
   }
@@ -136,7 +150,7 @@ export async function outputSample(taskId: number, sampleIds: number[], activeTx
   a.click();
 }
 
-export async function outputSamples(taskId: number, activeTxt: string) {
+export async function outputSamples(taskId: number, activeTxt: ExportType) {
   const samplesRes = await getSamples({ task_id: taskId, pageNo: 1, pageSize: 100000 });
   const sampleIdArrays = samplesRes.data;
   const sampleIds = [];

diff --git a/apps/frontend/src/api/types.ts b/apps/frontend/src/api/types.ts
@@ -88,7 +88,10 @@ export enum ExportType {
   COCO = 'COCO',
   YOLO = 'YOLO',
   CSV = 'CSV',
+  XML = 'XML',
   LABEL_ME = 'LABEL_ME',
+  TF_RECORD = 'TF_RECORD',
+  PASCAL_VOC = 'PASCAL_VOC',
 }
 
 export interface GetApiV1TasksTaskIdGetParams {

diff --git a/apps/frontend/src/components/ExportPortal/index.tsx b/apps/frontend/src/components/ExportPortal/index.tsx
@@ -1,5 +1,4 @@
-import type { RadioChangeEvent } from 'antd';
-import { Modal, Radio } from 'antd';
+import { Modal, Select } from 'antd';
 import React, { useCallback, useMemo, useState } from 'react';
 import { FlexLayout } from '@labelu/components-react';
 
@@ -18,17 +17,24 @@ export interface ExportPortalProps {
 export const exportDescriptionMapping = {
   [ExportType.JSON]: 'Label U 标准格式，包含任务id、标注结果、url、fileName字段',
   [ExportType.CSV]: '适用于单一标注类型的任务场景',
+  [ExportType.XML]: 'Label U标准格式的XML形式',
   [ExportType.COCO]: 'COCO数据集标准格式，面向物体检测（拉框）和图像分割（多边形）任务',
   [ExportType.MASK]: '面向图像分割（多边形）任务',
   [ExportType.YOLO]: 'YOLO数据集标准格式，面向物体检测（拉框）任务',
   [ExportType.LABEL_ME]: '兼容 Labelme 标注工具的标注数据格式（不支持立体框、曲线）',
+  [ExportType.TF_RECORD]: 'Tensorflow Object Detection API 标准格式',
+  [ExportType.PASCAL_VOC]: 'PASCAL VOC 标准格式，面向物体检测（拉框）任务',
 };
 
 const optionMapping = {
   [ExportType.JSON]: {
     label: ExportType.JSON,
     value: ExportType.JSON,
   },
+  [ExportType.XML]: {
+    label: ExportType.XML,
+    value: ExportType.XML,
+  },
   [ExportType.CSV]: {
     label: ExportType.CSV,
     value: ExportType.CSV,
@@ -41,6 +47,10 @@ const optionMapping = {
     label: ExportType.COCO,
     value: ExportType.COCO,
   },
+  [ExportType.PASCAL_VOC]: {
+    label: ExportType.PASCAL_VOC,
+    value: ExportType.PASCAL_VOC,
+  },
   [ExportType.MASK]: {
     label: ExportType.MASK,
     value: ExportType.MASK,
@@ -49,6 +59,10 @@ const optionMapping = {
     label: 'Labelme' as any,
     value: ExportType.LABEL_ME,
   },
+  [ExportType.TF_RECORD]: {
+    label: 'TF Record' as any,
+    value: ExportType.TF_RECORD,
+  },
 };
 
 function isIncludeCoco(tools?: any[]) {
@@ -74,7 +88,7 @@ export default function ExportPortal({ taskId, sampleIds, mediaType, tools, chil
     setModalVisible(false);
   }, []);
 
-  const handleOptionChange = ({ target: { value } }: RadioChangeEvent) => {
+  const handleOptionChange = (value: ExportType) => {
     setExportType(value);
   };
 
@@ -111,7 +125,7 @@ export default function ExportPortal({ taskId, sampleIds, mediaType, tools, chil
   }, [children, handleOpenModal]);
 
   const options = useMemo(() => {
-    const result = [optionMapping[ExportType.JSON]];
+    const result = [optionMapping[ExportType.JSON], optionMapping[ExportType.XML]];
 
     if (!mediaType) {
       return result;
@@ -124,6 +138,8 @@ export default function ExportPortal({ taskId, sampleIds, mediaType, tools, chil
     const onlyLineTool = tools?.length === 1 && tools[0].tool === 'lineTool';
 
     if (mediaType === MediaType.IMAGE) {
+      result.push(optionMapping[ExportType.TF_RECORD]);
+
       if (onlyPolygonTool || onlyRectTool || onlyPointTool || onlyCuboidTool || onlyLineTool) {
         result.push(optionMapping[ExportType.CSV]);
       }
@@ -133,7 +149,7 @@ export default function ExportPortal({ taskId, sampleIds, mediaType, tools, chil
       }
 
       if (onlyRectTool) {
-        result.push(optionMapping[ExportType.YOLO]);
+        result.push(optionMapping[ExportType.YOLO], optionMapping[ExportType.PASCAL_VOC]);
       }
 
       // mask: polygon
@@ -155,14 +171,8 @@ export default function ExportPortal({ taskId, sampleIds, mediaType, tools, chil
       <Modal title="选择导出格式" okText={'导出'} onOk={handleExport} onCancel={handleCloseModal} open={modalVisible}>
         <FlexLayout flex="column" gap="1rem">
           <FlexLayout.Header items="center" gap="1rem" flex>
-            <span>导出格式</span>
-            <Radio.Group
-              options={options}
-              onChange={handleOptionChange}
-              value={exportType}
-              optionType="button"
-              buttonStyle="solid"
-            />
+            <span style={{ whiteSpace: 'nowrap' }}>导出格式</span>
+            <Select popupMatchSelectWidth={false} options={options} onChange={handleOptionChange} value={exportType} />
           </FlexLayout.Header>
           <div>{exportDescriptionMapping[exportType]}</div>
         </FlexLayout>