Support data import for new tasks

backend/data_import/datasets.py

@@ -21,9 +21,12 @@ from .pipeline.readers import (
 )
 from label_types.models import CategoryType, LabelType, RelationType, SpanType
 from projects.models import (
+    BOUNDING_BOX,
     DOCUMENT_CLASSIFICATION,
+    IMAGE_CAPTIONING,
     IMAGE_CLASSIFICATION,
     INTENT_DETECTION_AND_SLOT_FILLING,
+    SEGMENTATION,
     SEQ2SEQ,
     SEQUENCE_LABELING,
     SPEECH2TEXT,
@@ -218,6 +221,9 @@ def select_dataset(project: Project, task: str, file_format: Format) -> Type[Dat
         SEQ2SEQ: Seq2seqDataset,
         INTENT_DETECTION_AND_SLOT_FILLING: CategoryAndSpanDataset,
         IMAGE_CLASSIFICATION: BinaryDataset,
+        IMAGE_CAPTIONING: BinaryDataset,
+        BOUNDING_BOX: BinaryDataset,
+        SEGMENTATION: BinaryDataset,
         SPEECH2TEXT: BinaryDataset,
     }
     if task not in mapping:

backend/data_import/pipeline/catalog.py

@@ -8,9 +8,12 @@ from typing_extensions import Literal
 
 from .exceptions import FileFormatException
 from projects.models import (
+    BOUNDING_BOX,
     DOCUMENT_CLASSIFICATION,
+    IMAGE_CAPTIONING,
     IMAGE_CLASSIFICATION,
     INTENT_DETECTION_AND_SLOT_FILLING,
+    SEGMENTATION,
     SEQ2SEQ,
     SEQUENCE_LABELING,
     SPEECH2TEXT,
@@ -432,16 +435,18 @@ Options.register(
     )
 )
 
-# Image Classification
-Options.register(
-    Option(
-        display_name=ImageFile.name,
-        task_id=IMAGE_CLASSIFICATION,
-        file_format=ImageFile,
-        arg=ArgNone,
-        file=IMAGE_CLASSIFICATION_DIR / "image_files.txt",
+# Image tasks
+image_tasks = [IMAGE_CLASSIFICATION, IMAGE_CAPTIONING, BOUNDING_BOX, SEGMENTATION]
+for task_name in image_tasks:
+    Options.register(
+        Option(
+            display_name=ImageFile.name,
+            task_id=task_name,
+            file_format=ImageFile,
+            arg=ArgNone,
+            file=IMAGE_CLASSIFICATION_DIR / "image_files.txt",
+        )
     )
-)
 
 # Speech to Text
 Options.register(

frontend/pages/projects/_id/dataset/index.vue

@@ -119,7 +119,8 @@ export default Vue.extend({
       return this.$route.params.id
     },
     isImageTask(): boolean {
-      return this.project.projectType === 'ImageClassification'
+      const imageTasks = ['ImageClassification', 'ImageCaptioning', 'BoundingBox', 'Segmentation']
+      return imageTasks.includes(this.project.projectType)
     },
     isAudioTask(): boolean {
       return this.project.projectType === 'Speech2text'