Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 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"""Image processor class for SegGPT."""16 17from typing import Optional, Union18 19import numpy as np20 21from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict22from ...image_transforms import resize, to_channel_dimension_format23from ...image_utils import (24 IMAGENET_DEFAULT_MEAN,25 IMAGENET_DEFAULT_STD,26 ChannelDimension,27 ImageInput,28 PILImageResampling,29 infer_channel_dimension_format,30 is_scaled_image,31 make_flat_list_of_images,32 to_numpy_array,33 valid_images,34)35from ...utils import TensorType, is_torch_available, logging, requires_backends36 37 38if is_torch_available():39 import torch40 41 42logger = logging.get_logger(__name__)43 44 45# See https://huggingface.co/papers/2212.02499 at 3.1 Redefining Output Spaces as "Images" - Semantic Segmentation from PAINTER paper46# Taken from https://github.com/Abdullah-Meda/Painter/blob/main/Painter/data/coco_semseg/gen_color_coco_panoptic_segm.py#L3147def build_palette(num_labels: int) -> list[tuple[int, int]]:48 base = int(num_labels ** (1 / 3)) + 149 margin = 256 // base50 51 # we assume that class_idx 0 is the background which is mapped to black52 color_list = [(0, 0, 0)]53 for location in range(num_labels):54 num_seq_r = location // base**255 num_seq_g = (location % base**2) // base56 num_seq_b = location % base57 58 R = 255 - num_seq_r * margin59 G = 255 - num_seq_g * margin60 B = 255 - num_seq_b * margin61 62 color_list.append((R, G, B))63 64 return color_list65 66 67def mask_to_rgb(68 mask: np.ndarray, palette: Optional[list[tuple[int, int]]] = None, data_format: Optional[ChannelDimension] = None69) -> np.ndarray:70 data_format = data_format if data_format is not None else ChannelDimension.FIRST71 72 if palette is not None:73 height, width = mask.shape74 75 rgb_mask = np.zeros((3, height, width), dtype=np.uint8)76 77 classes_in_mask = np.unique(mask)78 79 for class_idx in classes_in_mask:80 rgb_value = palette[class_idx]81 class_mask = (mask == class_idx).astype(np.uint8)82 class_mask = np.expand_dims(class_mask, axis=-1)83 class_rgb_mask = class_mask * np.array(rgb_value)84 class_rgb_mask = np.moveaxis(class_rgb_mask, -1, 0)85 rgb_mask += class_rgb_mask.astype(np.uint8)86 87 rgb_mask = np.clip(rgb_mask, 0, 255).astype(np.uint8)88 89 else:90 rgb_mask = np.repeat(mask[None, ...], 3, axis=0)91 92 return to_channel_dimension_format(rgb_mask, data_format)93 94 95class SegGptImageProcessor(BaseImageProcessor):96 r"""97 Constructs a SegGpt image processor.98 99 Args:100 do_resize (`bool`, *optional*, defaults to `True`):101 Whether to resize the image's (height, width) dimensions to the specified `(size["height"],102 size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.103 size (`dict`, *optional*, defaults to `{"height": 448, "width": 448}`):104 Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`105 method.106 resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):107 Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the108 `preprocess` method.109 do_rescale (`bool`, *optional*, defaults to `True`):110 Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`111 parameter in the `preprocess` method.112 rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):113 Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the114 `preprocess` method.115 do_normalize (`bool`, *optional*, defaults to `True`):116 Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`117 method.118 image_mean (`float` or `list[float]`, *optional*, defaults to `IMAGENET_DEFAULT_MEAN`):119 Mean to use if normalizing the image. This is a float or list of floats the length of the number of120 channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.121 image_std (`float` or `list[float]`, *optional*, defaults to `IMAGENET_DEFAULT_STD`):122 Standard deviation to use if normalizing the image. This is a float or list of floats the length of the123 number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.124 do_convert_rgb (`bool`, *optional*, defaults to `True`):125 Whether to convert the prompt mask to RGB format. Can be overridden by the `do_convert_rgb` parameter in the126 `preprocess` method.127 """128 129 model_input_names = ["pixel_values"]130 131 def __init__(132 self,133 do_resize: bool = True,134 size: Optional[dict[str, int]] = None,135 resample: PILImageResampling = PILImageResampling.BICUBIC,136 do_rescale: bool = True,137 rescale_factor: Union[int, float] = 1 / 255,138 do_normalize: bool = True,139 image_mean: Optional[Union[float, list[float]]] = None,140 image_std: Optional[Union[float, list[float]]] = None,141 do_convert_rgb: bool = True,142 **kwargs,143 ) -> None:144 super().__init__(**kwargs)145 size = size if size is not None else {"height": 448, "width": 448}146 size = get_size_dict(size)147 self.do_resize = do_resize148 self.do_rescale = do_rescale149 self.do_normalize = do_normalize150 self.size = size151 self.resample = resample152 self.rescale_factor = rescale_factor153 self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN154 self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD155 self.do_convert_rgb = do_convert_rgb156 157 def get_palette(self, num_labels: int) -> list[tuple[int, int]]:158 """Build a palette to map the prompt mask from a single channel to a 3 channel RGB.159 160 Args:161 num_labels (`int`):162 Number of classes in the segmentation task (excluding the background).163 164 Returns:165 `list[tuple[int, int]]`: Palette to map the prompt mask from a single channel to a 3 channel RGB.166 """167 return build_palette(num_labels)168 169 def mask_to_rgb(170 self,171 image: np.ndarray,172 palette: Optional[list[tuple[int, int]]] = None,173 data_format: Optional[Union[str, ChannelDimension]] = None,174 ) -> np.ndarray:175 """Converts a segmentation map to RGB format.176 177 Args:178 image (`np.ndarray`):179 Segmentation map with dimensions (height, width) where pixel values represent the class index.180 palette (`list[tuple[int, int]]`, *optional*, defaults to `None`):181 Palette to use to convert the mask to RGB format. If unset, the mask is duplicated across the channel182 dimension.183 data_format (`ChannelDimension` or `str`, *optional*):184 The channel dimension format for the output image. If unset, the channel dimension format of the input185 image is used. Can be one of:186 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.187 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.188 189 Returns:190 `np.ndarray`: The mask in RGB format.191 """192 return mask_to_rgb(image, palette=palette, data_format=data_format)193 194 # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize with PILImageResampling.BILINEAR->PILImageResampling.BICUBIC195 def resize(196 self,197 image: np.ndarray,198 size: dict[str, int],199 resample: PILImageResampling = PILImageResampling.BICUBIC,200 data_format: Optional[Union[str, ChannelDimension]] = None,201 input_data_format: Optional[Union[str, ChannelDimension]] = None,202 **kwargs,203 ) -> np.ndarray:204 """205 Resize an image to `(size["height"], size["width"])`.206 207 Args:208 image (`np.ndarray`):209 Image to resize.210 size (`dict[str, int]`):211 Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.212 resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):213 `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BICUBIC`.214 data_format (`ChannelDimension` or `str`, *optional*):215 The channel dimension format for the output image. If unset, the channel dimension format of the input216 image is used. Can be one of:217 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.218 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.219 - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.220 input_data_format (`ChannelDimension` or `str`, *optional*):221 The channel dimension format for the input image. If unset, the channel dimension format is inferred222 from the input image. Can be one of:223 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.224 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.225 - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.226 227 Returns:228 `np.ndarray`: The resized image.229 """230 size = get_size_dict(size)231 if "height" not in size or "width" not in size:232 raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")233 output_size = (size["height"], size["width"])234 return resize(235 image,236 size=output_size,237 resample=resample,238 data_format=data_format,239 input_data_format=input_data_format,240 **kwargs,241 )242 243 def _preprocess_step(244 self,245 images: ImageInput,246 do_resize: Optional[bool] = None,247 size: Optional[dict[str, int]] = None,248 resample: Optional[PILImageResampling] = None,249 do_rescale: Optional[bool] = None,250 rescale_factor: Optional[float] = None,251 do_normalize: Optional[bool] = None,252 image_mean: Optional[Union[float, list[float]]] = None,253 image_std: Optional[Union[float, list[float]]] = None,254 data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,255 input_data_format: Optional[Union[str, ChannelDimension]] = None,256 do_convert_rgb: Optional[bool] = None,257 num_labels: Optional[int] = None,258 **kwargs,259 ):260 """261 Preprocess an image or batch of images.262 263 Args:264 images (`ImageInput`):265 Image to _preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If266 passing in images with pixel values between 0 and 1, set `do_rescale=False`.267 do_resize (`bool`, *optional*, defaults to `self.do_resize`):268 Whether to resize the image.269 size (`dict[str, int]`, *optional*, defaults to `self.size`):270 Dictionary in the format `{"height": h, "width": w}` specifying the size of the output image after271 resizing.272 resample (`PILImageResampling` filter, *optional*, defaults to `self.resample`):273 `PILImageResampling` filter to use if resizing the image e.g. `PILImageResampling.BICUBIC`. Only has274 an effect if `do_resize` is set to `True`.275 do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):276 Whether to rescale the image values between [0 - 1].277 rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):278 Rescale factor to rescale the image by if `do_rescale` is set to `True`.279 do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):280 Whether to normalize the image.281 image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):282 Image mean to use if `do_normalize` is set to `True`.283 image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):284 Image standard deviation to use if `do_normalize` is set to `True`.285 return_tensors (`str` or `TensorType`, *optional*):286 The type of tensors to return. Can be one of:287 - Unset: Return a list of `np.ndarray`.288 - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.289 - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.290 - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.291 - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.292 data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):293 The channel dimension format for the output image. Can be one of:294 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.295 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.296 - Unset: Use the channel dimension format of the input image.297 input_data_format (`ChannelDimension` or `str`, *optional*):298 The channel dimension format for the input image. If unset, the channel dimension format is inferred299 from the input image. Can be one of:300 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.301 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.302 - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.303 do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):304 Whether to convert the prompt mask to RGB format. If `num_labels` is specified, a palette will be built305 to map the prompt mask from a single channel to a 3 channel RGB. If unset, the prompt mask is duplicated306 across the channel dimension. Must be set to `False` if the prompt mask is already in RGB format.307 num_labels: (`int`, *optional*):308 Number of classes in the segmentation task (excluding the background). If specified, a palette will be309 built, assuming that class_idx 0 is the background, to map the prompt mask from a single class_idx310 channel to a 3 channel RGB. Not specifying this will result in the prompt mask either being passed311 through as is if it is already in RGB format or being duplicated across the channel dimension.312 """313 do_resize = do_resize if do_resize is not None else self.do_resize314 do_rescale = do_rescale if do_rescale is not None else self.do_rescale315 do_normalize = do_normalize if do_normalize is not None else self.do_normalize316 do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb317 resample = resample if resample is not None else self.resample318 rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor319 image_mean = image_mean if image_mean is not None else self.image_mean320 image_std = image_std if image_std is not None else self.image_std321 322 size = size if size is not None else self.size323 size_dict = get_size_dict(size)324 325 # If segmentation map is passed we expect 2D images326 images = make_flat_list_of_images(images, expected_ndims=2 if do_convert_rgb else 3)327 328 if not valid_images(images):329 raise ValueError(330 "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "331 "torch.Tensor, tf.Tensor or jax.ndarray."332 )333 334 if do_resize and size is None:335 raise ValueError("Size must be specified if do_resize is True.")336 337 if do_rescale and rescale_factor is None:338 raise ValueError("Rescale factor must be specified if do_rescale is True.")339 340 if do_normalize and (image_mean is None or image_std is None):341 raise ValueError("Image mean and std must be specified if do_normalize is True.")342 343 # All transformations expect numpy arrays.344 images = [to_numpy_array(image) for image in images]345 346 if do_rescale and is_scaled_image(images[0]):347 logger.warning_once(348 "It looks like you are trying to rescale already rescaled images. If the input"349 " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."350 )351 352 if input_data_format is None and not do_convert_rgb:353 # We assume that all images have the same channel dimension format.354 input_data_format = infer_channel_dimension_format(images[0])355 356 if do_convert_rgb:357 palette = self.get_palette(num_labels) if num_labels is not None else None358 # Since this is the input for the next transformations its format should be the same as the input_data_format359 images = [360 self.mask_to_rgb(image=image, palette=palette, data_format=ChannelDimension.FIRST) for image in images361 ]362 input_data_format = ChannelDimension.FIRST363 364 if do_resize:365 images = [366 self.resize(image=image, size=size_dict, resample=resample, input_data_format=input_data_format)367 for image in images368 ]369 370 if do_rescale:371 images = [372 self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)373 for image in images374 ]375 376 if do_normalize:377 images = [378 self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)379 for image in images380 ]381 382 images = [383 to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images384 ]385 386 return images387 388 def preprocess(389 self,390 images: Optional[ImageInput] = None,391 prompt_images: Optional[ImageInput] = None,392 prompt_masks: Optional[ImageInput] = None,393 do_resize: Optional[bool] = None,394 size: Optional[dict[str, int]] = None,395 resample: Optional[PILImageResampling] = None,396 do_rescale: Optional[bool] = None,397 rescale_factor: Optional[float] = None,398 do_normalize: Optional[bool] = None,399 image_mean: Optional[Union[float, list[float]]] = None,400 image_std: Optional[Union[float, list[float]]] = None,401 do_convert_rgb: Optional[bool] = None,402 num_labels: Optional[int] = None,403 return_tensors: Optional[Union[str, TensorType]] = None,404 data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,405 input_data_format: Optional[Union[str, ChannelDimension]] = None,406 **kwargs,407 ):408 """409 Preprocess an image or batch of images.410 411 Args:412 images (`ImageInput`):413 Image to _preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If414 passing in images with pixel values between 0 and 1, set `do_rescale=False`.415 prompt_images (`ImageInput`):416 Prompt image to _preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If417 passing in images with pixel values between 0 and 1, set `do_rescale=False`.418 prompt_masks (`ImageInput`):419 Prompt mask from prompt image to _preprocess that specify prompt_masks value in the preprocessed output.420 Can either be in the format of segmentation maps (no channels) or RGB images. If in the format of421 RGB images, `do_convert_rgb` should be set to `False`. If in the format of segmentation maps, `num_labels`422 specifying `num_labels` is recommended to build a palette to map the prompt mask from a single channel to423 a 3 channel RGB. If `num_labels` is not specified, the prompt mask will be duplicated across the channel424 dimension.425 do_resize (`bool`, *optional*, defaults to `self.do_resize`):426 Whether to resize the image.427 size (`dict[str, int]`, *optional*, defaults to `self.size`):428 Dictionary in the format `{"height": h, "width": w}` specifying the size of the output image after429 resizing.430 resample (`PILImageResampling` filter, *optional*, defaults to `self.resample`):431 `PILImageResampling` filter to use if resizing the image e.g. `PILImageResampling.BICUBIC`. Only has432 an effect if `do_resize` is set to `True`. Doesn't apply to prompt mask as it is resized using nearest.433 do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):434 Whether to rescale the image values between [0 - 1].435 rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):436 Rescale factor to rescale the image by if `do_rescale` is set to `True`.437 do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):438 Whether to normalize the image.439 image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):440 Image mean to use if `do_normalize` is set to `True`.441 image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):442 Image standard deviation to use if `do_normalize` is set to `True`.443 do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):444 Whether to convert the prompt mask to RGB format. If `num_labels` is specified, a palette will be built445 to map the prompt mask from a single channel to a 3 channel RGB. If unset, the prompt mask is duplicated446 across the channel dimension. Must be set to `False` if the prompt mask is already in RGB format.447 num_labels: (`int`, *optional*):448 Number of classes in the segmentation task (excluding the background). If specified, a palette will be449 built, assuming that class_idx 0 is the background, to map the prompt mask from a plain segmentation map450 with no channels to a 3 channel RGB. Not specifying this will result in the prompt mask either being passed451 through as is if it is already in RGB format (if `do_convert_rgb` is false) or being duplicated452 across the channel dimension.453 return_tensors (`str` or `TensorType`, *optional*):454 The type of tensors to return. Can be one of:455 - Unset: Return a list of `np.ndarray`.456 - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.457 - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.458 - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.459 - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.460 data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):461 The channel dimension format for the output image. Can be one of:462 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.463 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.464 - Unset: Use the channel dimension format of the input image.465 input_data_format (`ChannelDimension` or `str`, *optional*):466 The channel dimension format for the input image. If unset, the channel dimension format is inferred467 from the input image. Can be one of:468 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.469 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.470 - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.471 """472 if all(v is None for v in [images, prompt_images, prompt_masks]):473 raise ValueError("At least one of images, prompt_images, prompt_masks must be specified.")474 475 data = {}476 477 if images is not None:478 images = self._preprocess_step(479 images,480 is_mask=False,481 do_resize=do_resize,482 size=size,483 resample=resample,484 do_rescale=do_rescale,485 rescale_factor=rescale_factor,486 do_normalize=do_normalize,487 image_mean=image_mean,488 image_std=image_std,489 do_convert_rgb=False,490 data_format=data_format,491 input_data_format=input_data_format,492 **kwargs,493 )494 495 data["pixel_values"] = images496 497 if prompt_images is not None:498 prompt_images = self._preprocess_step(499 prompt_images,500 is_mask=False,501 do_resize=do_resize,502 size=size,503 resample=resample,504 do_rescale=do_rescale,505 rescale_factor=rescale_factor,506 do_normalize=do_normalize,507 image_mean=image_mean,508 image_std=image_std,509 do_convert_rgb=False,510 data_format=data_format,511 input_data_format=input_data_format,512 **kwargs,513 )514 515 data["prompt_pixel_values"] = prompt_images516 517 if prompt_masks is not None:518 prompt_masks = self._preprocess_step(519 prompt_masks,520 do_resize=do_resize,521 size=size,522 resample=PILImageResampling.NEAREST,523 do_rescale=do_rescale,524 rescale_factor=rescale_factor,525 do_normalize=do_normalize,526 image_mean=image_mean,527 image_std=image_std,528 do_convert_rgb=do_convert_rgb,529 num_labels=num_labels,530 data_format=data_format,531 input_data_format=input_data_format,532 **kwargs,533 )534 535 data["prompt_masks"] = prompt_masks536 537 return BatchFeature(data=data, tensor_type=return_tensors)538 539 def post_process_semantic_segmentation(540 self, outputs, target_sizes: Optional[list[tuple[int, int]]] = None, num_labels: Optional[int] = None541 ):542 """543 Converts the output of [`SegGptImageSegmentationOutput`] into segmentation maps. Only supports544 PyTorch.545 546 Args:547 outputs ([`SegGptImageSegmentationOutput`]):548 Raw outputs of the model.549 target_sizes (`list[tuple[int, int]]`, *optional*):550 List of length (batch_size), where each list item (`tuple[int, int]`) corresponds to the requested551 final size (height, width) of each prediction. If left to None, predictions will not be resized.552 num_labels (`int`, *optional*):553 Number of classes in the segmentation task (excluding the background). If specified, a palette will be554 built, assuming that class_idx 0 is the background, to map prediction masks from RGB values to class555 indices. This value should be the same used when preprocessing inputs.556 Returns:557 semantic_segmentation: `list[torch.Tensor]` of length `batch_size`, where each item is a semantic558 segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is559 specified). Each entry of each `torch.Tensor` correspond to a semantic class id.560 """561 requires_backends(self, ["torch"])562 # batch_size x num_channels x 2*height x width563 masks = outputs.pred_masks564 565 # Predicted mask and prompt are concatenated in the height dimension566 # batch_size x num_channels x height x width567 masks = masks[:, :, masks.shape[2] // 2 :, :]568 569 # To unnormalize we need to permute to channel last570 # batch_size x height x width x num_channels571 std = torch.tensor(self.image_std).to(masks.device)572 mean = torch.tensor(self.image_mean).to(masks.device)573 574 masks = masks.permute(0, 2, 3, 1) * std + mean575 576 # batch_size x num_channels x height x width577 masks = masks.permute(0, 3, 1, 2)578 579 # Clip to match with palette if specified580 masks = torch.clip(masks * 255, 0, 255)581 582 semantic_segmentation = []583 palette_tensor = None584 palette = self.get_palette(num_labels) if num_labels is not None else None585 if palette is not None:586 palette_tensor = torch.tensor(palette).to(device=masks.device, dtype=torch.float)587 _, num_channels, _, _ = masks.shape588 palette_tensor = palette_tensor.view(1, 1, num_labels + 1, num_channels)589 590 for idx, mask in enumerate(masks):591 if target_sizes is not None:592 mask = torch.nn.functional.interpolate(593 mask.unsqueeze(0),594 size=target_sizes[idx],595 mode="nearest",596 )[0]597 598 if num_labels is not None:599 channels, height, width = mask.shape600 dist = mask.permute(1, 2, 0).view(height, width, 1, channels)601 dist = dist - palette_tensor602 dist = torch.pow(dist, 2)603 dist = torch.sum(dist, dim=-1)604 pred = dist.argmin(dim=-1)605 606 else:607 # If no palette is specified SegGpt will try to paint using the mask class idx as RGB608 pred = mask.mean(dim=0).int()609 610 semantic_segmentation.append(pred)611 612 return semantic_segmentation613 614 615__all__ = ["SegGptImageProcessor"]616