riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 23 Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 18 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/pipelines/base.py
View File

@ -872,6 +872,9 @@ class Pipeline(_ScikitCompat):
if self.feature_extractor is not None:
self.feature_extractor.save_pretrained(save_directory)
if self.image_processor is not None:
self.image_processor.save_pretrained(save_directory)
if self.modelcard is not None:
self.modelcard.save_pretrained(save_directory)

15
tests/pipelines/test_pipelines_image_segmentation.py
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)