Save image_processor while saving pipeline (ImageSegmentationPipeline) (#25884)
* Save image_processor while saving pipeline (ImageSegmentationPipeline) * Fix black issues
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user