Save image_processor while saving pipeline (ImageSegmentationPipeline) (#25884)

* Save image_processor while saving pipeline (ImageSegmentationPipeline)

* Fix black issues
This commit is contained in:
raghavanone
2023-08-31 19:38:20 +05:30
committed by GitHub
parent a39ebbf879
commit 2be8a9098e
2 changed files with 18 additions and 0 deletions

View File

@@ -13,6 +13,7 @@
# limitations under the License.
import hashlib
import tempfile
import unittest
from typing import Dict
@@ -714,3 +715,17 @@ class ImageSegmentationPipelineTests(unittest.TestCase):
},
],
)
def test_save_load(self):
model_id = "hf-internal-testing/tiny-detr-mobilenetsv3-panoptic"
model = AutoModelForImageSegmentation.from_pretrained(model_id)
image_processor = AutoImageProcessor.from_pretrained(model_id)
image_segmenter = pipeline(
task="image-segmentation",
model=model,
image_processor=image_processor,
)
with tempfile.TemporaryDirectory() as tmpdirname:
image_segmenter.save_pretrained(tmpdirname)
pipeline(task="image-segmentation", model=tmpdirname)