Add flat images to existing task defs

Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>
caikit · Jan 22, 2024 · f967218 · f967218
1 parent a0ea2e2
commit f967218
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/caikit_computer_vision/data_model/__init__.py b/caikit_computer_vision/data_model/__init__.py
@@ -14,6 +14,7 @@
 
 # Local
 from . import image_classification, image_segmentation, object_detection, tasks
+from .flat_image import *
 from .image_classification import *
 from .image_segmentation import *
 from .object_detection import *
diff --git a/caikit_computer_vision/data_model/tasks.py b/caikit_computer_vision/data_model/tasks.py
@@ -22,14 +22,15 @@
 from caikit.core import TaskBase, task
 
 # Local
+from .flat_image import FlatImage
 from .image_classification import ImageClassificationResult
 from .image_segmentation import ImageSegmentationResult
 from .object_detection import ObjectDetectionResult
 
 
 # TODO - add support for image DM primitives
 @task(
-    required_parameters={"inputs": Union[bytes, str]},
+    required_parameters={"inputs": Union[bytes, str, FlatImage]},
     output_type=ObjectDetectionResult,
 )
 class ObjectDetectionTask(TaskBase):
@@ -40,7 +41,7 @@ class ObjectDetectionTask(TaskBase):
 
 
 @task(
-    required_parameters={"inputs": Union[bytes, str]},
+    required_parameters={"inputs": Union[bytes, str, FlatImage]},
     output_type=ImageClassificationResult,
 )
 class ImageClassificationTask(TaskBase):
@@ -51,7 +52,7 @@ class ImageClassificationTask(TaskBase):
 
 
 @task(
-    required_parameters={"inputs": Union[bytes, str]},
+    required_parameters={"inputs": Union[bytes, str, FlatImage]},
     output_type=ImageSegmentationResult,
 )
 class ImageSegmentationTask(TaskBase):