Add DPT support for ORTModelForSemanticSegmentation (#1759)

* add DPT support * make style
huggingface · Mar 18, 2024 · 7e08a82 · 7e08a82
1 parent 23fe6f7
commit 7e08a82
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -557,6 +557,8 @@ class TasksManager:
             "feature-extraction",
             "image-feature-extraction",
             "depth-estimation",
+            "image-segmentation",
+            "semantic-segmentation",
             onnx="DptOnnxConfig",
         ),
         "electra": supported_tasks_mapping(

diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -2961,7 +2961,7 @@ def test_compare_to_io_binding(self, model_arch):
 
 
 class ORTModelForSemanticSegmentationIntegrationTest(ORTModelTestMixin):
-    SUPPORTED_ARCHITECTURES = ("segformer",)
+    SUPPORTED_ARCHITECTURES = ("segformer", "dpt")
 
     FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES}
     ORTMODEL_CLASS = ORTModelForSemanticSegmentation

diff --git a/tests/onnxruntime/utils_onnxruntime_tests.py b/tests/onnxruntime/utils_onnxruntime_tests.py
@@ -87,6 +87,7 @@
     "deit": "hf-internal-testing/tiny-random-DeiTModel",
     "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder",
     "detr": "hf-internal-testing/tiny-random-detr",
+    "dpt": "hf-internal-testing/tiny-random-DPTModel",
     "distilbert": "hf-internal-testing/tiny-random-DistilBertModel",
     "electra": "hf-internal-testing/tiny-random-ElectraModel",
     "encoder-decoder": {