CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
modular_segformer.py152 linesDownload Raw Back to segformer
1# coding=utf-82# Copyright 2025 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Fast Image processor class for Segformer."""16 17from typing import Optional, Union18 19import torch20from torchvision.transforms.v2 import functional as F21 22from transformers.models.beit.image_processing_beit_fast import BeitFastImageProcessorKwargs, BeitImageProcessorFast23 24from ...image_processing_utils import BatchFeature25from ...image_processing_utils_fast import (26    group_images_by_shape,27    reorder_images,28)29from ...image_utils import (30    IMAGENET_DEFAULT_MEAN,31    IMAGENET_DEFAULT_STD,32    ChannelDimension,33    ImageInput,34    PILImageResampling,35    SizeDict,36)37from ...processing_utils import Unpack38from ...utils import (39    TensorType,40)41 42 43class SegformerFastImageProcessorKwargs(BeitFastImageProcessorKwargs):44    pass45 46 47class SegformerImageProcessorFast(BeitImageProcessorFast):48    resample = PILImageResampling.BILINEAR49    image_mean = IMAGENET_DEFAULT_MEAN50    image_std = IMAGENET_DEFAULT_STD51    size = {"height": 512, "width": 512}52    do_resize = True53    do_rescale = True54    rescale_factor = 1 / 25555    do_normalize = True56    do_reduce_labels = False57    do_center_crop = None58    crop_size = None59 60    def _preprocess_image_like_inputs(61        self,62        images: ImageInput,63        segmentation_maps: Optional[ImageInput],64        do_convert_rgb: bool,65        input_data_format: ChannelDimension,66        device: Optional[Union[str, "torch.device"]] = None,67        **kwargs: Unpack[SegformerFastImageProcessorKwargs],68    ) -> BatchFeature:69        """70        Preprocess image-like inputs.71        """72        images = self._prepare_image_like_inputs(73            images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device74        )75        images_kwargs = kwargs.copy()76        images_kwargs["do_reduce_labels"] = False77        batch_feature = self._preprocess(images, **images_kwargs)78 79        if segmentation_maps is not None:80            processed_segmentation_maps = self._prepare_image_like_inputs(81                images=segmentation_maps,82                expected_ndims=2,83                do_convert_rgb=False,84                input_data_format=ChannelDimension.FIRST,85            )86 87            segmentation_maps_kwargs = kwargs.copy()88            segmentation_maps_kwargs.update(89                {90                    "do_normalize": False,91                    "do_rescale": False,92                    # Nearest interpolation is used for segmentation maps instead of BILINEAR.93                    "interpolation": F.InterpolationMode.NEAREST_EXACT,94                }95            )96            processed_segmentation_maps = self._preprocess(97                images=processed_segmentation_maps, **segmentation_maps_kwargs98            ).pixel_values99            batch_feature["labels"] = processed_segmentation_maps.squeeze(1).to(torch.int64)100 101        return batch_feature102 103    def _preprocess(104        self,105        images: list["torch.Tensor"],106        do_reduce_labels: bool,107        interpolation: Optional["F.InterpolationMode"],108        do_resize: bool,109        do_rescale: bool,110        do_normalize: bool,111        size: SizeDict,112        rescale_factor: float,113        image_mean: Union[float, list[float]],114        image_std: Union[float, list[float]],115        disable_grouping: bool,116        return_tensors: Optional[Union[str, TensorType]],117        **kwargs,118    ) -> BatchFeature:  # Return type can be list if return_tensors=None119        if do_reduce_labels:120            images = self.reduce_label(images)  # Apply reduction if needed121 122        # Group images by size for batched resizing123        resized_images = images124        if do_resize:125            grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)126            resized_images_grouped = {}127            for shape, stacked_images in grouped_images.items():128                resized_stacked_images = self.resize(image=stacked_images, size=size, interpolation=interpolation)129                resized_images_grouped[shape] = resized_stacked_images130            resized_images = reorder_images(resized_images_grouped, grouped_images_index)131 132        # Group images by size for further processing (rescale/normalize)133        # Needed in case do_resize is False, or resize returns images with different sizes134        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)135        processed_images_grouped = {}136        for shape, stacked_images in grouped_images.items():137            # Fused rescale and normalize138            stacked_images = self.rescale_and_normalize(139                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std140            )141            processed_images_grouped[shape] = stacked_images142 143        processed_images = reorder_images(processed_images_grouped, grouped_images_index)144 145        # Stack images into a single tensor if return_tensors is set146        processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images147 148        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)149 150 151__all__ = ["SegformerImageProcessorFast"]152 
Aluode/PerceptionLabPortable · CoolFace