# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pathlib
from dataclasses import dataclass
from typing import Any

import torch
import torch.nn.functional as F
from huggingface_hub.dataclasses import strict
from torch import Tensor, nn
from torchvision.transforms.v2 import functional as tvF

from ...activations import ACT2FN
from ...backbone_utils import BackboneConfigMixin, consolidate_backbone_kwargs_to_config
from ...configuration_utils import PreTrainedConfig
from ...image_processing_utils import BatchFeature
from ...image_transforms import (
    center_to_corners_format,
)
from ...image_utils import (
    AnnotationFormat,
    AnnotationType,
    ChannelDimension,
    PILImageResampling,
    SizeDict,
    get_max_height_width,
    validate_annotations,
)
from ...modeling_outputs import BackboneOutput, BaseModelOutput
from ...processing_utils import Unpack
from ...utils import (
    TensorType,
    auto_docstring,
    logging,
    torch_int,
)
from ...utils.generic import ModelOutput, TransformersKwargs, can_return_tuple
from ..clip.modeling_clip import CLIPMLP
from ..convnext.modeling_convnext import ConvNextLayer
from ..detr.image_processing_detr import (
    SUPPORTED_ANNOTATION_FORMATS,
    DetrImageProcessor,
    convert_segmentation_to_rle,
    prepare_coco_detection_annotation,
)
from ..dinov2.configuration_dinov2 import Dinov2Config
from ..dinov2.modeling_dinov2 import (
    Dinov2Backbone,
    Dinov2Embeddings,
    Dinov2Encoder,
    Dinov2Layer,
    Dinov2PreTrainedModel,
)
from ..lw_detr.configuration_lw_detr import LwDetrConfig
from ..lw_detr.modeling_lw_detr import (
    LwDetrC2FLayer,
    LwDetrConvEncoder,
    LwDetrConvNormLayer,
    LwDetrForObjectDetection,
    LwDetrLayerNorm,
    LwDetrModel,
    LwDetrObjectDetectionOutput,
    LwDetrPreTrainedModel,
    refine_bboxes,
)


logger = logging.get_logger(__name__)


class RfDetrImageProcessor(DetrImageProcessor):
    def _preprocess(
        self,
        images: list["torch.Tensor"],
        annotations: AnnotationType | list[AnnotationType] | None,
        return_segmentation_masks: bool,
        masks_path: str | pathlib.Path | None,
        do_resize: bool,
        size: SizeDict,
        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        do_convert_annotations: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        do_pad: bool,
        pad_size: SizeDict | None,
        format: str | AnnotationFormat | None,
        return_tensors: str | TensorType | None,
        **kwargs,
    ) -> BatchFeature:
        """
        Preprocess an image or a batch of images so that it can be used by the model.
        """
        if annotations is not None and isinstance(annotations, dict):
            annotations = [annotations]

        if annotations is not None and len(images) != len(annotations):
            raise ValueError(
                f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match."
            )

        format = AnnotationFormat(format)
        if annotations is not None:
            validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations)

        if (
            masks_path is not None
            and format == AnnotationFormat.COCO_PANOPTIC
            and not isinstance(masks_path, (pathlib.Path, str))
        ):
            raise ValueError(
                "The path to the directory containing the mask PNG files should be provided as a"
                f" `pathlib.Path` or string object, but is {type(masks_path)} instead."
            )

        data = {}

        processed_images = []
        processed_annotations = []
        pixel_masks = []  # Initialize pixel_masks here
        for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)):
            # prepare (COCO annotations as a list of Dict -> DETR target as a single Dict per image)
            if annotations is not None:
                annotation = self.prepare_annotation(
                    image,
                    annotation,
                    format,
                    return_segmentation_masks=return_segmentation_masks,
                    masks_path=masks_path,
                    input_data_format=ChannelDimension.FIRST,
                )
            # Rescale then resize like in the original RF-DETR implementation
            if do_rescale:
                image = self.rescale(image, rescale_factor)
            if do_resize:
                resized_image = self.resize(image, size=size, resample=resample)
                if annotations is not None:
                    annotation = self.resize_annotation(
                        annotation,
                        orig_size=image.size()[-2:],
                        target_size=resized_image.size()[-2:],
                    )
                image = resized_image
            if do_normalize:
                image = self.normalize(image, image_mean, image_std)
            if do_convert_annotations and annotations is not None:
                annotation = self.normalize_annotation(annotation, image.shape[-2:])

            processed_images.append(image)
            processed_annotations.append(annotation)
        images = processed_images
        annotations = processed_annotations if annotations is not None else None

        if do_pad:
            # depends on all resized image shapes so we need another loop
            if pad_size is not None:
                padded_size = (pad_size.height, pad_size.width)
            else:
                padded_size = get_max_height_width(images)

            padded_images = []
            padded_annotations = []
            for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)):
                # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...}
                if padded_size == image.size()[-2:]:
                    padded_images.append(image)
                    pixel_masks.append(torch.ones(padded_size, dtype=torch.int64, device=image.device))
                    padded_annotations.append(annotation)
                    continue
                image, pixel_mask, annotation = self.pad(
                    image, padded_size, annotation=annotation, update_bboxes=do_convert_annotations
                )
                padded_images.append(image)
                padded_annotations.append(annotation)
                pixel_masks.append(pixel_mask)
            images = padded_images
            annotations = padded_annotations if annotations is not None else None
            data.update({"pixel_mask": torch.stack(pixel_masks, dim=0)})

        data.update({"pixel_values": torch.stack(images, dim=0)})
        encoded_inputs = BatchFeature(data, tensor_type=return_tensors)
        if annotations is not None:
            encoded_inputs["labels"] = [
                BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations
            ]
        return encoded_inputs

    def prepare_annotation(
        self,
        image: torch.Tensor,
        target: dict,
        format: AnnotationFormat | None = None,
        return_segmentation_masks: bool | None = None,
        masks_path: str | pathlib.Path | None = None,
        input_data_format: str | ChannelDimension | None = None,
    ) -> dict:
        """
        Prepare an annotation for feeding into RF_DETR model.
        """
        format = format if format is not None else self.format

        if format == AnnotationFormat.COCO_DETECTION:
            return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks
            target = prepare_coco_detection_annotation(
                image, target, return_segmentation_masks, input_data_format=input_data_format
            )
        else:
            raise ValueError(f"Format {format} is not supported.")
        return target

    def post_process_object_detection(
        self,
        outputs,
        threshold: float = 0.5,
        target_sizes: TensorType | list[tuple] | None = None,
        top_k: int | None = None,
    ):
        out_logits, out_bbox = outputs.logits, outputs.pred_boxes
        if target_sizes is not None:
            if len(out_logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )
        top_k = top_k if top_k is not None else out_logits.shape[1]
        prob = out_logits.sigmoid()
        batch_size, num_queries, num_classes = prob.shape
        k_value = min(top_k, num_queries * num_classes)
        scores, topk_indexes = torch.topk(prob.view(batch_size, -1), k_value, dim=1)
        topk_boxes = topk_indexes // num_classes
        labels = topk_indexes % num_classes
        boxes = center_to_corners_format(out_bbox)
        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))

        if target_sizes is not None:
            if isinstance(target_sizes, list):
                img_h = torch.tensor([i[0] for i in target_sizes], device=boxes.device)
                img_w = torch.tensor([i[1] for i in target_sizes], device=boxes.device)
            else:
                img_h, img_w = target_sizes.to(boxes.device).unbind(1)
            scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device)
            boxes = boxes * scale_fct[:, None, :]

        results = []
        for score, label, box in zip(scores, labels, boxes):
            keep = score > threshold
            results.append({"scores": score[keep], "labels": label[keep], "boxes": box[keep]})
        return results

    def post_process_instance_segmentation(
        self,
        outputs,
        threshold: float = 0.5,
        mask_threshold: float = 0.0,
        target_sizes: list[tuple[int, int]] | None = None,
        return_coco_annotation: bool = False,
        return_binary_maps: bool = False,
        top_k: int | None = None,
        **kwargs,
    ) -> list[dict[str, Any]]:
        """
        Converts the output of [`RfDetrForInstanceSegmentation`] into instance segmentation predictions.

        Args:
            outputs ([`RfDetrInstanceSegmentationOutput`]):
                Raw outputs of the model.
            threshold (`float`, *optional*, defaults to 0.5):
                Score threshold to keep predicted instance masks.
            mask_threshold (`float`, *optional*, defaults to 0.0):
                Threshold to binarize predicted masks.
            target_sizes (`list[tuple[int, int]]`, *optional*):
                Target ``(height, width)`` for each image. If unset, masks are not resized.
            return_coco_annotation (`bool`, *optional*, defaults to `False`):
                If `True`, return segmentation maps as COCO run-length encoding instead of tensors.
                Mutually exclusive with `return_binary_maps`.
            return_binary_maps (`bool`, *optional*, defaults to `False`):
                If `True`, return segmentation maps as a stacked tensor of binary instance masks
                (one per detected instance), without overlap resolution. This matches the output
                format of the original ``rfdetr`` package. Mutually exclusive with
                `return_coco_annotation`.
            top_k (`int`, *optional*):
                Maximum number of candidate queries evaluated before score thresholding.
                Defaults to the total number of queries.

        Returns:
            `list[dict]`: One dict per image with keys:
            - **segmentation** -- `Tensor[H, W]` of ``int32`` segment ids (``-1`` = background),
              `Tensor[num_instances, H, W]` of bool binary masks when `return_binary_maps=True`,
              or a list of RLE encodings when `return_coco_annotation=True`.
            - **segments_info** -- List of dicts with keys ``id``, ``label_id``, and ``score``.
        """
        if return_coco_annotation and return_binary_maps:
            raise ValueError("`return_coco_annotation` and `return_binary_maps` cannot both be `True`.")
        out_logits, out_masks = outputs.logits, outputs.pred_masks
        top_k = top_k if top_k is not None else out_logits.shape[1]

        prob = out_logits.sigmoid()
        batch_size, num_queries, num_classes = prob.shape
        top_k = min(top_k, num_queries * num_classes)

        scores, topk_indexes = torch.topk(prob.view(batch_size, -1), top_k, dim=1)
        topk_queries = topk_indexes // num_classes
        labels = topk_indexes % num_classes
        masks = torch.gather(
            out_masks,
            1,
            topk_queries.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, out_masks.shape[-2], out_masks.shape[-1]),
        )

        results: list[dict[str, Any]] = []
        for batch_idx, (batch_scores, batch_labels, batch_masks) in enumerate(zip(scores, labels, masks)):
            keep = batch_scores > threshold
            pred_scores = batch_scores[keep]
            pred_labels = batch_labels[keep]
            mask_logits = batch_masks[keep]

            if target_sizes is not None:
                height, width = int(target_sizes[batch_idx][0]), int(target_sizes[batch_idx][1])
            else:
                height, width = mask_logits.shape[-2], mask_logits.shape[-1]

            if pred_scores.shape[0] == 0:
                segmentation = torch.full((height, width), -1, dtype=torch.int32, device=out_logits.device)
                results.append({"segmentation": segmentation, "segments_info": []})
                continue

            mask_logits_resized = F.interpolate(
                mask_logits.unsqueeze(1).float(),
                size=(height, width),
                mode="bilinear",
                align_corners=False,
            ).squeeze(1)

            segmentation = torch.full((height, width), -1, dtype=torch.int32, device=out_logits.device)
            segments: list[dict] = []
            instance_maps: list = []
            current_id = 0

            for i in range(pred_scores.shape[0]):
                binary_mask = mask_logits_resized[i] > mask_threshold
                if not binary_mask.any():
                    continue
                if return_binary_maps:
                    instance_maps.append(binary_mask)
                else:
                    pixels_to_paint = binary_mask & (segmentation == -1)
                    if not pixels_to_paint.any():
                        continue
                current_id += 1
                if not return_binary_maps:
                    segmentation[pixels_to_paint] = current_id
                segments.append(
                    {"id": current_id, "label_id": int(pred_labels[i]), "score": round(pred_scores[i].item(), 6)}
                )

            if return_coco_annotation:
                segmentation = convert_segmentation_to_rle(segmentation)
            elif return_binary_maps:
                segmentation = (
                    torch.stack(instance_maps, dim=0)
                    if instance_maps
                    else torch.zeros(0, height, width, dtype=torch.bool, device=out_logits.device)
                )

            results.append({"segmentation": segmentation, "segments_info": segments})

        return results

    def post_process_panoptic_segmentation(self):
        raise NotImplementedError("Panoptic segmentation is not supported for RF-DETR.")

    def post_process_semantic_segmentation(self):
        raise NotImplementedError("Semantic segmentation is not supported for RF-DETR.")


@auto_docstring(checkpoint="Roboflow/rf-detr-base")
@strict
class RfDetrDinov2Config(Dinov2Config):
    r"""
    layerscale_value (`float`, *optional*, defaults to 1.0):
        Initial value to use for layer scale.
    drop_path_rate (`float`, *optional*, defaults to 0.0):
        Stochastic depth rate per sample (when applied in the main path of residual layers).
    use_swiglu_ffn (`bool`, *optional*, defaults to `False`):
        Whether to use the SwiGLU feedforward neural network.
    apply_layernorm (`bool`, *optional*, defaults to `True`):
        Whether to apply layer normalization to the feature maps in case the model is used as backbone.
    reshape_hidden_states (`bool`, *optional*, defaults to `True`):
        Whether to reshape the feature maps to 4D tensors of shape `(batch_size, d_model, height, width)` in
        case the model is used as backbone. If `False`, the feature maps will be 3D tensors of shape `(batch_size,
        seq_len, d_model)`.
    use_mask_token (`bool`, *optional*, defaults to `True`):
        Whether to use mask_token in embeddings.
    num_windows (`int`, *optional*, defaults to 4):
        Number of windows to use for windowed attention. If 1, no windowed attention is used.

    Example:

    ```python
    >>> from transformers import RfDetrDinov2Config, RfDetrDinov2Backbone

    >>> # Initializing a RfDetrDinov2 base style configuration
    >>> configuration = RfDetrDinov2Config()

    >>> # Initializing a model (with random weights) from the base style configuration
    >>> model = RfDetrDinov2Backbone(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```"""

    model_type = "rf_detr_dinov2"

    num_windows: int = 4

    def __post_init__(self, **kwargs):
        self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, self.num_hidden_layers + 1)]
        self.set_output_features_output_indices(
            out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
        )
        window_block_indexes = set(range(self._out_indices[-1] + 1))
        window_block_indexes.difference_update(self._out_indices)
        self.window_block_indexes = list(window_block_indexes)
        BackboneConfigMixin.__post_init__(**kwargs)


@auto_docstring(checkpoint="Roboflow/rf-detr-base")
@strict
class RfDetrConfig(LwDetrConfig):
    r"""
    hidden_expansion (`float`, *optional*, defaults to 0.5):
        Expansion factor for hidden dimensions in the projector layers.
    c2f_num_blocks (`int`, *optional*, defaults to 3):
        Number of blocks in the C2F layer.
    activation_function (`str`, *optional*, defaults to `"silu"`):
        The non-linear activation function in the projector. Supported values are `"silu"`, `"relu"`, `"gelu"`.
    decoder_n_points (`int`, *optional*, defaults to 4):
        The number of sampled keys in each feature level for each attention head in the decoder.
    decoder_layers (`int`, *optional*, defaults to 3):
        Number of decoder layers in the transformer.
    decoder_self_attention_heads (`int`, *optional*, defaults to 8):
        Number of attention heads for each attention layer in the decoder self-attention.
    decoder_cross_attention_heads (`int`, *optional*, defaults to 16):
        Number of attention heads for each attention layer in the decoder cross-attention.
    decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
        The non-linear activation function in the decoder. Supported values are `"relu"`, `"silu"`, `"gelu"`.
    num_queries (`int`, *optional*, defaults to 300):
        Number of object queries, i.e. detection slots. This is the maximal number of objects
        [`RfDetrModel`] can detect in a single image.
    group_detr (`int`, *optional*, defaults to 13):
        Number of groups for Group DETR attention mechanism, which helps reduce computational complexity.
    disable_custom_kernels (`bool`, *optional*, defaults to `True`):
        Disable the use of custom CUDA and CPU kernels. This option is necessary for the ONNX export, as custom
        kernels are not supported by PyTorch ONNX export.
    class_loss_coefficient (`float`, *optional*, defaults to 1):
        Relative weight of the classification loss in the Hungarian matching cost.
    dice_loss_coefficient (`float`, *optional*, defaults to 1):
        Relative weight of the DICE/F-1 loss in the object detection loss.
    bbox_loss_coefficient (`float`, *optional*, defaults to 5):
        Relative weight of the L1 bounding box loss in the object detection loss.
    giou_loss_coefficient (`float`, *optional*, defaults to 2):
        Relative weight of the generalized IoU loss in the object detection loss.
    num_feature_levels (`int`, *optional*, defaults to 1):
        Number of feature levels used in the multiscale deformable attention.
    mask_loss_coefficient (`float`, *optional*, defaults to 1):
        Relative weight of the Focal loss in the instance segmentation mask loss.
    mask_point_sample_ratio (`int`, *optional*, defaults to 16):
        The ratio of points to sample for the mask loss calculation.
    mask_downsample_ratio (`int`, *optional*, defaults to 4):
        The downsample ratio for the segmentation masks compared to the input image resolution.
    mask_class_loss_coefficient (`float`, *optional*, defaults to 5.0):
        Relative weight of the Focal loss in the instance segmentation loss.
    mask_dice_loss_coefficient (`float`, *optional*, defaults to 5.0):
        Relative weight of the DICE/F-1 loss in the instance segmentation loss.
    segmentation_head_activation_function (`str`, *optional*, defaults to `"gelu"`):
        The non-linear activation function in the segmentation head. Supported values are `"relu"`, `"silu"`, `"gelu"`.

    Examples:

    ```python
    >>> from transformers import RfDetrConfig, RfDetrModel

    >>> # Initializing a RF-DETR roboflow/rf-detr-base style configuration
    >>> configuration = RfDetrConfig()

    >>> # Initializing a model (with random weights) from the Roboflow/rf-detr-base style configuration
    >>> model = RfDetrModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```"""

    model_type = "rf_detr"

    layer_norm_eps: float = 1e-5
    dropout: float = 0.1
    num_feature_levels: int = 1
    mask_loss_coefficient: int | float = 1
    mask_point_sample_ratio: int = 16
    mask_downsample_ratio: int = 4
    mask_class_loss_coefficient: int | float = 5.0
    mask_dice_loss_coefficient: int | float = 5.0
    segmentation_head_activation_function: str = "gelu"
    intermediate_size: int = 1024

    batch_norm_eps = AttributeError("batch_norm_eps is replaced by layer_norm_eps in RfDetrConfig.")
    projector_scale_factors = AttributeError("RfDetr only uses a single scale layer instead of multiple scale layers.")

    def __post_init__(self, **kwargs):
        self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
            backbone_config=self.backbone_config,
            default_config_type="rf_detr_dinov2",
            default_config_kwargs={
                "num_attention_heads": 6,
                "out_features": ["stage2", "stage5", "stage8", "stage11"],
                "hidden_size": 384,
                "num_register_tokens": 0,
                "image_size": 518,
            },
            **kwargs,
        )
        PreTrainedConfig.__post_init__(**kwargs)

    def validate_architecture(self):
        raise AttributeError(
            "validate_architecture is not used in RfDetrConfig because it does not rely on multiple scale layers."
        )


class RfDetrDinov2Embeddings(Dinov2Embeddings):
    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
        """
        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
        images. This method is also adapted to support torch.jit tracing and interpolation at torch.float32 precision.

        Adapted from:
        - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
        - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
        """

        num_patches = embeddings.shape[1] - 1
        num_positions = self.position_embeddings.shape[1] - 1

        # always interpolate when tracing to ensure the exported model works for dynamic input shapes
        if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
            return self.position_embeddings

        class_pos_embed = self.position_embeddings[:, :1]
        patch_pos_embed = self.position_embeddings[:, 1:]

        dim = embeddings.shape[-1]

        new_height = height // self.patch_size
        new_width = width // self.patch_size

        sqrt_num_positions = torch_int(num_positions**0.5)
        patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
        target_dtype = patch_pos_embed.dtype
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.to(torch.float32),
            size=(new_height, new_width),
            mode="bicubic",
            # Difference from Dinov2, we use align_corners=False and antialias=True
            align_corners=False,
            antialias=True,
        ).to(dtype=target_dtype)

        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)

        return torch.cat((class_pos_embed, patch_pos_embed), dim=1)

    def window_partition(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
        """
        Splits each image's patch-token grid into num_windows^2 local windows,
        replicates the [CLS] token per window, and returns window-local token sequences
        """
        batch_size = embeddings.shape[0]
        num_windows = self.config.num_windows
        patch_size = self.patch_size
        num_height_patches = height // patch_size
        num_width_patches = width // patch_size
        num_width_patches_per_window = num_width_patches // num_windows
        num_height_patches_per_window = num_height_patches // num_windows

        # Split the embeddings into the [CLS] token and the pixel tokens
        cls_token_with_pos_embed = embeddings[:, :1]
        pixel_tokens_with_pos_embed = embeddings[:, 1:]
        pixel_tokens_with_pos_embed = pixel_tokens_with_pos_embed.view(
            batch_size, num_height_patches, num_width_patches, -1
        )

        # Reshape the pixel tokens into windowed pixel tokens
        windowed_pixel_tokens = pixel_tokens_with_pos_embed.view(
            batch_size, num_windows, num_width_patches_per_window, num_windows, num_height_patches_per_window, -1
        )
        windowed_pixel_tokens = windowed_pixel_tokens.transpose(2, 3)
        windowed_pixel_tokens = windowed_pixel_tokens.reshape(
            batch_size * num_windows**2, num_height_patches_per_window * num_width_patches_per_window, -1
        )

        # Repeat the [CLS] token per window
        windowed_cls_token_with_pos_embed = cls_token_with_pos_embed.repeat(num_windows**2, 1, 1)

        # Concatenate the [CLS] token with the windowed pixel tokens to get the final embeddings
        embeddings = torch.cat((windowed_cls_token_with_pos_embed, windowed_pixel_tokens), dim=1)
        return embeddings

    def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None) -> torch.Tensor:
        batch_size, _, height, width = pixel_values.shape
        target_dtype = self.patch_embeddings.projection.weight.dtype
        embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))

        if bool_masked_pos is not None and self.use_mask_token:
            embeddings = torch.where(
                bool_masked_pos.unsqueeze(-1), self.mask_token.to(embeddings.dtype).unsqueeze(0), embeddings
            )

        # add the [CLS] token to the embedded patch tokens
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        # add positional encoding to each token
        embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)

        # Difference from Dinov2, we use window partitioning
        if self.config.num_windows > 1:
            embeddings = self.window_partition(embeddings, height, width)
        embeddings = self.dropout(embeddings)

        return embeddings


class RfDetrDinov2Layer(Dinov2Layer):
    def __init__(self, config: RfDetrDinov2Config, layer_idx: int):
        super().__init__(config)
        self.num_windows = config.num_windows
        self.global_attention = layer_idx not in config.window_block_indexes

    def window_unpartition_before_attention(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        For layers configured to use global attention, merges the window-batched sequences back
        into one sequence per image so attention can be computed across all windows jointly.
        """
        batch_size, seq_len, channels = hidden_states.shape
        num_windows_squared = self.num_windows**2
        hidden_states = hidden_states.view(batch_size // num_windows_squared, num_windows_squared * seq_len, channels)
        return hidden_states

    def window_partition_after_attention(
        self, hidden_state_shape: tuple[int, int, int], self_attention_output: torch.Tensor
    ) -> torch.Tensor:
        """
        After global attention, reshapes the output sequence back into window-batched
        form so the model can continue in the same windowed pipeline.
        """

        batch_size, seq_len, channels = hidden_state_shape
        num_windows_squared = self.num_windows**2
        self_attention_output = self_attention_output.view(
            batch_size * num_windows_squared, seq_len // num_windows_squared, channels
        )
        return self_attention_output

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        residual = hidden_states

        # Difference from Dinov2, when the layer is not a window block, we need to unpartition the hidden states before the attention
        if self.global_attention:
            hidden_states = self.window_unpartition_before_attention(hidden_states)

        hidden_states_norm = self.norm1(hidden_states)
        self_attention_output = self.attention(hidden_states_norm)

        # And reverse the operation after the attention
        if self.global_attention:
            self_attention_output = self.window_partition_after_attention(hidden_states.shape, self_attention_output)

        self_attention_output = self.layer_scale1(self_attention_output)

        # first residual connection
        hidden_states = self.drop_path(self_attention_output) + residual
        residual = hidden_states

        # in Dinov2, layernorm is also applied after self-attention
        hidden_states = self.norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.layer_scale2(hidden_states)

        # second residual connection
        hidden_states = self.drop_path(hidden_states) + residual

        return hidden_states


class RfDetrDinov2PreTrainedModel(Dinov2PreTrainedModel):
    pass


class RfDetrDinov2Encoder(Dinov2Encoder):
    def __init__(self, config: RfDetrDinov2Config):
        super().__init__(config)
        self.layer = nn.ModuleList([RfDetrDinov2Layer(config, i) for i in range(config.num_hidden_layers)])


class RfDetrDinov2Backbone(Dinov2Backbone):
    def window_unpartition(self, hidden_state: torch.Tensor, height: int, width: int) -> torch.Tensor:
        """
        Reassembles windowed patch tokens into their original 2D patch layout (image-level grid structure)
        before converting backbone hidden states into spatial feature maps.
        """
        num_windows = self.config.num_windows
        patch_size = self.config.patch_size
        num_h_patches = height // patch_size
        num_w_patches = width // patch_size
        hidden_batch_size, seq_len, channels = hidden_state.shape
        num_windows_squared = num_windows**2
        num_h_patches_per_window = num_h_patches // num_windows
        num_w_patches_per_window = num_w_patches // num_windows

        # Reshape the hidden states into the original sequence length
        hidden_state = hidden_state.reshape(
            hidden_batch_size // num_windows_squared, num_windows_squared * seq_len, channels
        )
        hidden_state = hidden_state.view(
            hidden_batch_size // num_windows_squared,
            num_windows,
            num_windows,
            num_h_patches_per_window,
            num_w_patches_per_window,
            channels,
        )
        hidden_state = hidden_state.transpose(2, 3)
        return hidden_state

    def forward(
        self,
        pixel_values: torch.Tensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BackboneOutput:
        r"""
        Examples:

        ```python
        >>> from transformers import AutoImageProcessor, AutoBackbone
        >>> import torch
        >>> from PIL import Image
        >>> import requests

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
        >>> model = AutoBackbone.from_pretrained(
        ...     "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"]
        ... )

        >>> inputs = processor(image, return_tensors="pt")

        >>> outputs = model(**inputs)
        >>> feature_maps = outputs.feature_maps
        >>> list(feature_maps[-1].shape)
        [1, 768, 16, 16]
        ```"""
        # Like Dinov2, we need to output the hidden states to extract the layers for the stages
        kwargs["output_hidden_states"] = True

        embedding_output = self.embeddings(pixel_values)
        output: BaseModelOutput = self.encoder(embedding_output, **kwargs)
        hidden_states = output.hidden_states

        feature_maps = ()
        for stage, hidden_state in zip(self.stage_names, hidden_states):
            if stage in self.out_features:
                if self.config.apply_layernorm:
                    hidden_state = self.layernorm(hidden_state)
                if self.config.reshape_hidden_states:
                    hidden_state = hidden_state[:, 1:]
                    # this was actually a bug in the original implementation that we copied here,
                    # cause normally the order is height, width
                    batch_size, _, height, width = pixel_values.shape
                    num_h_patches = height // self.config.patch_size
                    num_w_patches = width // self.config.patch_size

                    # Difference from Dinov2, when the layer is not a window block, we need to unpartition the hidden states before reshaping
                    if self.config.num_windows > 1:
                        hidden_state = self.window_unpartition(hidden_state, height, width)

                    hidden_state = hidden_state.reshape(batch_size, num_h_patches, num_w_patches, -1)
                    hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()

                feature_maps += (hidden_state,)

        return BackboneOutput(
            feature_maps=tuple(feature_maps),
            hidden_states=hidden_states,
            attentions=output.attentions,
        )


class RfDetrLayerNorm(LwDetrLayerNorm):
    pass


class RfDetrConvNormLayer(LwDetrConvNormLayer):
    def __init__(
        self,
        config: RfDetrConfig,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int,
        activation: str | None = None,
    ):
        super().__init__(
            config,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            activation,
        )
        self.norm = RfDetrLayerNorm(out_channels, data_format="channels_first", eps=config.layer_norm_eps)


class RfDetrC2FLayer(LwDetrC2FLayer):
    pass


class RfDetrScaleProjector(nn.Module):
    def __init__(self, config: RfDetrConfig):
        super().__init__()
        projector_input_dim: int = config.backbone_config.hidden_size * len(config.backbone_config.out_indices)
        self.projector_layer = RfDetrC2FLayer(config, projector_input_dim)
        self.layer_norm = RfDetrLayerNorm(config.d_model, data_format="channels_first")

    def forward(self, hidden_states: tuple[torch.Tensor]) -> torch.Tensor:
        hidden_states = torch.cat(hidden_states, dim=1)
        hidden_states = self.projector_layer(hidden_states)
        hidden_states = self.layer_norm(hidden_states)
        return hidden_states


class RfDetrConvEncoder(LwDetrConvEncoder):
    def __init__(self, config: RfDetrConfig):
        super().__init__(config)
        self.backbone = RfDetrDinov2Backbone(config.backbone_config)
        self.projector = RfDetrScaleProjector(config)

    def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
        # send pixel_values through the model to get list of feature maps
        features = self.backbone(pixel_values).feature_maps
        features = self.projector(features)
        mask = nn.functional.interpolate(pixel_mask[None].float(), size=features.shape[-2:]).to(torch.bool)[0]
        return features, mask


class RfDetrPreTrainedModel(LwDetrPreTrainedModel):
    # Roboflow checkpoints use bare keys with no top-level prefix
    _checkpoint_conversion_prefix_free = True

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        if hasattr(module, "segmentation_bias") and isinstance(module.segmentation_bias, nn.Parameter):
            nn.init.constant_(module.segmentation_bias, 0.0)


@auto_docstring(
    custom_intro="""
    Base class for outputs of the RfDetr backbone-decoder model.
    """
)
@dataclass
class RfDetrModelOutput(ModelOutput):
    r"""
    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
        Initial reference points sent through the Transformer decoder.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, d_model)`):
        Stacked intermediate hidden states (output of each layer of the decoder).
    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
        Stacked intermediate reference points (reference points of each layer of the decoder).
    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
        foreground and background).
    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
        Logits of predicted bounding boxes coordinates in the first stage.
    backbone_features (list of `torch.FloatTensor` of shape `(batch_size, config.num_channels, config.image_size, config.image_size)`):
        Features from the backbone.
    """

    last_hidden_state: torch.FloatTensor | None = None
    init_reference_points: torch.FloatTensor | None = None
    intermediate_hidden_states: torch.FloatTensor | None = None
    intermediate_reference_points: torch.FloatTensor | None = None
    enc_outputs_class: torch.FloatTensor | None = None
    enc_outputs_coord_logits: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None
    cross_attentions: tuple[torch.FloatTensor, ...] | None = None
    backbone_features: list[torch.Tensor] = None


class RfDetrModel(LwDetrModel):
    def __init__(self, config: RfDetrConfig):
        super().__init__(config)
        self.d_model = config.d_model

    def generate_topk_proposals(
        self, group_id: int, object_query_embedding: Tensor, output_proposals: Tensor, invalid_mask: Tensor, topk: int
    ) -> tuple[Tensor, Tensor, Tensor]:
        """
        Generates and selects the top-k object query embeddings and bounding box proposals for a specific query group.

        The pipeline proceeds as follows:

        1. Project and normalize the base query embeddings for the specific query group.
        2. Predict classification scores and bounding box refinements for the current query features.
        3. Apply the predicted deltas to the initial proposals to obtain refined spatial coordinates.
        4. Identify the indices of the highest-confidence predictions and gather the refined coordinates for these
           top-k candidates (detached to prevent gradient flow back to the proposal generation stage).
        5. Gather the associated query features to be used as starting points for the decoder stage.
        """
        # Step 1.
        object_query = self.enc_output[group_id](object_query_embedding)
        object_query = self.enc_output_norm[group_id](object_query)

        # Step 2.
        enc_outputs_class_proposals = self.enc_out_class_embed[group_id](object_query)
        delta_bbox = self.enc_out_bbox_embed[group_id](object_query)
        enc_outputs_class_proposals = enc_outputs_class_proposals.masked_fill(
            invalid_mask.to(enc_outputs_class_proposals.device), float("-inf")
        )

        # Step 3.
        enc_outputs_coord = refine_bboxes(output_proposals, delta_bbox)

        # Step 4.
        topk_proposals = torch.topk(enc_outputs_class_proposals.max(-1)[0], topk, dim=1)[1]
        topk_coords_logits_undetach = torch.gather(
            enc_outputs_coord,
            1,
            topk_proposals.unsqueeze(-1).expand(-1, -1, 4),
        )
        topk_coords_logits = topk_coords_logits_undetach.detach()

        # Step 5.
        object_query_undetach = torch.gather(
            object_query, 1, topk_proposals.unsqueeze(-1).expand(-1, -1, self.config.d_model)
        )
        return object_query_undetach, topk_coords_logits, topk_coords_logits_undetach

    @can_return_tuple
    @auto_docstring(
        custom_intro="""
    Forward pass of the RF-DETR model. The pipeline proceeds as follows:

        1. Generate an initial set of object query embeddings and spatial location proposals from the
           backbone's flattened output.
        2. Initialize storage for refined encoder-stage predictions (accommodating multi-group query
           structures) and iteratively refine object queries and their coordinates for each query group
           to capture the highest-confidence candidates from the encoder stage.
        3. Initialize learnable query features and spatial reference points (restricting to the primary
           group during inference for efficiency).
        4. Project the base reference points across the batch, refine them with the predicted coordinate
           refinements (shifting attention to the discovered object locations before decoding), and expand
           the target query features to match the batch dimensions.
        5. Pass the refined queries and updated reference points through the transformer decoder to
            aggregate detailed spatial context from the multi-scale features.
    """
    )
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        pixel_mask: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> RfDetrModelOutput:
        r"""
        Examples:

        ```python
        >>> from transformers import AutoImageProcessor, RfDetrModel
        >>> from PIL import Image
        >>> import requests

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> image_processor = AutoImageProcessor.from_pretrained("Roboflow/rf-detr-base")
        >>> model = RfDetrModel.from_pretrained("Roboflow/rf-detr-base")

        >>> inputs = image_processor(images=image, return_tensors="pt")

        >>> outputs = model(**inputs)

        >>> last_hidden_states = outputs.last_hidden_state
        >>> list(last_hidden_states.shape)
        [1, 200, 256]
        ```"""
        batch_size, num_channels, height, width = pixel_values.shape
        device = pixel_values.device

        if pixel_mask is None:
            pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)

        features, mask = self.backbone(pixel_values, pixel_mask)

        source_flatten = features.flatten(2).transpose(1, 2)
        mask_flatten = mask.flatten(1)
        spatial_shapes_list = [features.shape[2:]]
        spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=source_flatten.device)
        level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
        valid_ratios = self.get_valid_ratio(mask, dtype=source_flatten.dtype).unsqueeze(1)

        # Step 1.
        object_query_embedding, output_proposals, invalid_mask = self.gen_encoder_output_proposals(
            source_flatten, ~mask_flatten, spatial_shapes_list
        )

        # Step 2.
        group_detr = self.group_detr if self.training else 1
        topk = self.num_queries
        topk_coords_logits = torch.empty(
            (batch_size, topk * group_detr, 4), device=self.device, dtype=output_proposals.dtype
        )
        enc_outputs_coord_logits = torch.empty(
            (batch_size, topk * group_detr, 4), device=self.device, dtype=output_proposals.dtype
        )
        enc_outputs_class = torch.empty(
            (batch_size, topk * group_detr, self.config.d_model), device=self.device, dtype=output_proposals.dtype
        )
        for group_id in range(group_detr):
            object_query_undetach, group_topk_coords_logits, topk_coords_logits_undetach = (
                self.generate_topk_proposals(group_id, object_query_embedding, output_proposals, invalid_mask, topk)
            )
            topk_coords_logits[:, group_id * topk : (group_id + 1) * topk] = group_topk_coords_logits
            enc_outputs_coord_logits[:, group_id * topk : (group_id + 1) * topk] = topk_coords_logits_undetach
            enc_outputs_class[:, group_id * topk : (group_id + 1) * topk] = object_query_undetach

        # Step 3.
        if self.training:
            reference_points = self.reference_point_embed.weight
            query_feat = self.query_feat.weight
        else:
            reference_points = self.reference_point_embed.weight[: self.num_queries]
            query_feat = self.query_feat.weight[: self.num_queries]

        # Step 4.
        reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1)
        two_stage_len = enc_outputs_coord_logits.shape[-2]
        reference_points_two_stage_subset = reference_points[..., :two_stage_len, :]
        reference_points_subset = reference_points[..., two_stage_len:, :]
        reference_points_two_stage_subset = refine_bboxes(topk_coords_logits, reference_points_two_stage_subset)
        reference_points = torch.cat([reference_points_two_stage_subset, reference_points_subset], dim=-2)
        init_reference_points = reference_points
        target = query_feat.unsqueeze(0).expand(batch_size, -1, -1)

        # Step 5.
        decoder_outputs = self.decoder(
            inputs_embeds=target,
            reference_points=reference_points,
            spatial_shapes=spatial_shapes,
            spatial_shapes_list=spatial_shapes_list,
            level_start_index=level_start_index,
            valid_ratios=valid_ratios,
            encoder_hidden_states=source_flatten,
            encoder_attention_mask=mask_flatten,
            **kwargs,
        )

        return RfDetrModelOutput(
            init_reference_points=init_reference_points,
            last_hidden_state=decoder_outputs.last_hidden_state,
            intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
            intermediate_reference_points=decoder_outputs.intermediate_reference_points,
            backbone_features=features,
            enc_outputs_class=enc_outputs_class,
            enc_outputs_coord_logits=enc_outputs_coord_logits,
            hidden_states=decoder_outputs.hidden_states,
            attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
        )


class RfDetrObjectDetectionOutput(LwDetrObjectDetectionOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
        Total loss as a linear combination of a negative log-likelihood (cross-entropy) for class prediction and a
        bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
        scale-invariant IoU loss.
    loss_dict (`Dict`, *optional*):
        A dictionary containing the individual losses. Useful for logging.
    logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
        Classification logits (including no-object) for all queries.
    pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
        Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
        values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
        possible padding). You can use [`~DeformableDetrProcessor.post_process_object_detection`] to retrieve the
        unnormalized bounding boxes.
    auxiliary_outputs (`list[Dict]`, *optional*):
        Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
        and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
        `pred_boxes`) for each decoder layer.
    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
        Initial reference points sent through the Transformer decoder.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, d_model)`):
        Stacked intermediate hidden states (output of each layer of the decoder).
    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
        Stacked intermediate reference points (reference points of each layer of the decoder).
    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
        foreground and background).
    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
        Logits of predicted bounding boxes coordinates in the first stage.
    backbone_features (list of `torch.FloatTensor` of shape `(batch_size, config.num_channels, config.image_size, config.image_size)`):
        Features from the backbone.
    """

    backbone_features: list[torch.Tensor] = None


class RfDetrForObjectDetection(LwDetrForObjectDetection):
    def predict_encoder_class_logits(self, enc_outputs_class: torch.Tensor) -> Tensor:
        """
        Predicts classification logits from encoder hidden states for each query group.
        """
        enc_outputs_class_list = enc_outputs_class.split(self.config.num_queries, dim=1)
        group_detr = self.config.group_detr if self.training else 1
        pred_class = [
            self.model.enc_out_class_embed[group_index](enc_outputs_class_list[group_index])
            for group_index in range(group_detr)
        ]
        return torch.cat(pred_class, dim=1)

    def predict_class_and_boxes(
        self, hidden_states: torch.Tensor, reference_points: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Predicts classification logits and refined bounding boxes from transformer hidden states and reference points.
        """
        logits = self.class_embed(hidden_states)
        boxes_delta = self.bbox_embed(hidden_states)
        boxes = refine_bboxes(reference_points, boxes_delta)
        return logits, boxes

    @can_return_tuple
    @auto_docstring(
        custom_intro="""
    The forward pass proceeds as follows:

        1. Process the visual input through the base RF-DETR model to obtain the transformer's last hidden state and
           the final sequence of reference points.
        2. First stage: Generate classification logits from the encoder's proposed object query embeddings.
        3. Second stage: Predict the final classification labels and refined bounding boxes using the decoder's last hidden state
           and the most recent reference points.

    """
    )
    def forward(
        self,
        pixel_values: torch.FloatTensor = None,
        pixel_mask: torch.LongTensor | None = None,
        labels: list[dict] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> RfDetrObjectDetectionOutput:
        r"""
        labels (`list[Dict]` of len `(batch_size,)`, *optional*):
            Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
            following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
            respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
            in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.

        Examples:

        ```python
        >>> from transformers import AutoImageProcessor, RfDetrForObjectDetection
        >>> from PIL import Image
        >>> import requests

        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> image_processor = AutoImageProcessor.from_pretrained("Roboflow/rf-detr-base")
        >>> model = RfDetrForObjectDetection.from_pretrained("Roboflow/rf-detr-base")

        >>> inputs = image_processor(images=image, return_tensors="pt")
        >>> outputs = model(**inputs)

        >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
        >>> target_sizes = torch.tensor([image.size[::-1]])
        >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
        ...     0
        ... ]
        >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        ...     box = [round(i, 2) for i in box.tolist()]
        ...     print(
        ...         f"Detected {model.config.id2label[label.item()]} with confidence "
        ...         f"{round(score.item(), 3)} at location {box}"
        ...     )
        Detected cat with confidence 0.8 at location [16.5, 52.84, 318.25, 470.78]
        Detected cat with confidence 0.789 at location [342.19, 24.3, 640.02, 372.25]
        Detected remote with confidence 0.633 at location [40.79, 72.78, 176.76, 117.25]
        ```"""
        # Step 1.
        outputs = self.model(pixel_values, pixel_mask=pixel_mask, **kwargs)

        last_hidden_states = outputs.last_hidden_state
        intermediate_reference_points = outputs.intermediate_reference_points

        # Step 2.
        enc_outputs_class_logits = self.predict_encoder_class_logits(outputs.enc_outputs_class)

        # Step 3.
        logits, pred_boxes = self.predict_class_and_boxes(last_hidden_states, intermediate_reference_points[-1])

        loss, loss_dict, auxiliary_outputs = None, None, None
        if labels is not None:
            outputs_class, outputs_coord = None, None
            if self.config.auxiliary_loss:
                outputs_class, outputs_coord = self.predict_class_and_boxes(
                    outputs.intermediate_hidden_states, intermediate_reference_points
                )

            loss, loss_dict, auxiliary_outputs = self.loss_function(
                logits,
                labels,
                self.device,
                pred_boxes,
                self.config,
                outputs_class,
                outputs_coord,
                enc_outputs_class_logits,
                outputs.enc_outputs_coord_logits,
            )

        return RfDetrObjectDetectionOutput(
            loss=loss,
            loss_dict=loss_dict,
            logits=logits,
            pred_boxes=pred_boxes,
            auxiliary_outputs=auxiliary_outputs,
            last_hidden_state=outputs.last_hidden_state,
            intermediate_hidden_states=outputs.intermediate_hidden_states,
            intermediate_reference_points=outputs.intermediate_reference_points,
            init_reference_points=outputs.init_reference_points,
            enc_outputs_class=enc_outputs_class_logits,
            enc_outputs_coord_logits=outputs.enc_outputs_coord_logits,
            backbone_features=outputs.backbone_features,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )


@auto_docstring(
    custom_intro="""
    Output type of [`RfDetrForInstanceSegmentation`].
    """
)
@dataclass
class RfDetrInstanceSegmentationOutput(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
        Total loss as a linear combination of a negative log-likelihood (cross-entropy) for class prediction and a
        bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
        scale-invariant IoU loss.
    loss_dict (`Dict`, *optional*):
        A dictionary containing the individual losses. Useful for logging.
    logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
        Classification logits (including no-object) for all queries.
    pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
        Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
        values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
        possible padding). You can use [`~DeformableDetrProcessor.post_process_object_detection`] to retrieve the
        unnormalized bounding boxes.
    pred_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height/4, width/4)`):
        Segmentation masks logits for all queries. See also
        [`~RfDetrImageProcessor.post_process_instance_segmentation`] to obtain instance segmentation maps.
    auxiliary_outputs (`list[Dict]`, *optional*):
        Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
        and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
        `pred_boxes`) for each decoder layer.
    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
        Initial reference points sent through the Transformer decoder.
    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_queries, d_model)`, *optional*):
        Sequence of hidden-states at the output of the last layer of the decoder of the model.
    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, d_model)`):
        Stacked intermediate hidden states (output of each layer of the decoder).
    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
        Stacked intermediate reference points (reference points of each layer of the decoder).
    enc_outputs_mask_logits (`torch.FloatTensor` of shape `(batch_size, num_queries, width, height)`, *optional*):
        Mask logits from the encoder for all queries.
    """

    loss: torch.FloatTensor | None = None
    loss_dict: dict | None = None
    logits: torch.FloatTensor | None = None
    pred_boxes: torch.FloatTensor | None = None
    pred_masks: torch.FloatTensor = None
    auxiliary_outputs: list[dict] | None = None
    init_reference_points: torch.FloatTensor | None = None
    last_hidden_state: torch.FloatTensor | None = None
    intermediate_hidden_states: torch.FloatTensor | None = None
    intermediate_reference_points: torch.FloatTensor | None = None
    enc_outputs_mask_logits: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None
    cross_attentions: tuple[torch.FloatTensor, ...] | None = None


class RfDetrSegmentationBlock(ConvNextLayer):
    def __init__(self, config: RfDetrConfig):
        dim = config.d_model
        super().__init__(config)
        self.depthwise_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim)
        self.layernorm = RfDetrLayerNorm(dim, eps=1e-6)
        self.pointwise_conv = nn.Linear(dim, dim)
        self.act = ACT2FN[config.segmentation_head_activation_function]
        del self.dwconv
        del self.pwconv1
        del self.pwconv2
        del self.layer_scale_parameter
        del self.drop_path

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.depthwise_conv(hidden_states)
        hidden_states = hidden_states.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        hidden_states = self.layernorm(hidden_states)
        hidden_states = self.pointwise_conv(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = hidden_states.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
        hidden_states = hidden_states + residual
        return hidden_states


class RfDetrSegmentationMLP(CLIPMLP):
    def __init__(self, config: RfDetrConfig):
        super().__init__(config)
        self.activation_fn = ACT2FN[config.segmentation_head_activation_function]
        self.fc1 = nn.Linear(config.d_model, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.d_model)


class RfDetrSegmentationMLPBlock(nn.Module):
    def __init__(self, config: RfDetrConfig):
        super().__init__()
        dim = config.d_model
        self.norm = nn.LayerNorm(dim)
        self.mlp = RfDetrSegmentationMLP(config)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.norm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = hidden_states + residual
        return hidden_states


class RfDetrForInstanceSegmentation(RfDetrPreTrainedModel):
    # When using clones, all layers > 0 will be clones, but layer 0 *is* required
    # We can't initialize the model on meta device as some weights are modified during the initialization
    _no_split_modules = None

    def __init__(self, config: RfDetrConfig):
        super().__init__(config)

        self.model = RfDetrForObjectDetection(config)

        num_blocks = config.decoder_layers
        self.downsample_ratio = config.mask_downsample_ratio
        self.blocks = nn.ModuleList([RfDetrSegmentationBlock(config) for _ in range(num_blocks)])
        self.spatial_features_proj = nn.Conv2d(config.d_model, config.d_model, kernel_size=1)

        self.query_features_block = RfDetrSegmentationMLPBlock(config)
        self.query_features_proj = nn.Linear(config.d_model, config.d_model)

        self.segmentation_bias = nn.Parameter(torch.zeros(1), requires_grad=True)

        self.post_init()

    def get_mask_logits(self, query_features: Tensor, spatial_features: Tensor) -> Tensor:
        """
        Compute the per-query mask logits.
        Query features are projected to the same dimension as the spatial features and then multiplied
        with the spatial features to get the mask logits.
        The mask logits are then reshaped to the spatial dimensions and broadcast with a segmentation bias
        parameter.

        Args:
            query_features (`torch.Tensor`): Query features of shape (batch_size, num_queries, d_model).
            spatial_features (`torch.Tensor`): Spatial features of shape (batch_size, hidden_dim, height, width).

        Returns:
            `torch.Tensor`: Mask logits of shape (batch_size, num_queries, height, width).
        """
        batch_size, num_queries, _ = query_features.shape
        height, width = spatial_features.shape[2], spatial_features.shape[3]

        query_features = self.query_features_block(query_features)
        query_features = self.query_features_proj(query_features)
        mask_logits = torch.matmul(query_features, spatial_features.flatten(2))
        mask_logits = mask_logits.view(batch_size, num_queries, height, width)
        mask_logits = mask_logits + self.segmentation_bias
        return mask_logits

    def segmentation_head(
        self, spatial_features, list_query_features, image_size: torch.Size, skip_blocks: bool = False
    ) -> list[torch.Tensor] | torch.Tensor:
        """
        Compute mask logits from spatial features and query features.

        Args:
            spatial_features: Multi-scale spatial features of shape
                (batch_size, num_channels, feature_height, feature_width).
            list_query_features: When `skip_blocks` is False, a list of query feature tensors of shape
                (batch_size, num_queries, d_model) for each decoder layer. When `skip_blocks` is True,
                a single tensor of shape (batch_size, num_queries, d_model).
            image_size: Original image spatial dimensions (height, width).
            skip_blocks: If True, skip the convolutional blocks and compute mask logits directly.

        Returns:
            When `skip_blocks` is False: list of mask logit tensors of shape
            (batch_size, num_queries, mask_height, mask_width), where mask size is image size divided
            by `downsample_ratio`. When `skip_blocks` is True: a single such tensor.
        """
        target_size = (image_size[0] // self.downsample_ratio, image_size[1] // self.downsample_ratio)
        spatial_features = F.interpolate(spatial_features, size=target_size, mode="bilinear", align_corners=False)

        if not skip_blocks:
            list_mask_logits = []
            for block, query_features in zip(self.blocks, list_query_features):
                spatial_features = block(spatial_features)
                spatial_features_proj = self.spatial_features_proj(spatial_features)
                mask_logits = self.get_mask_logits(query_features, spatial_features_proj)
                list_mask_logits.append(mask_logits)
        else:
            list_mask_logits = self.get_mask_logits(list_query_features, spatial_features)

        return list_mask_logits

    @can_return_tuple
    @auto_docstring(
        custom_intro="""
    Forward pass of the RF-DETR model for instance segmentation. The pipeline proceeds as follows:

        1. Process the visual input through the base RF-DETR model to obtain multi-scale spatial features,
           query embeddings, and their transformation history.
        2. Generate classification logits and initial segmentation masks from the encoder's proposed
           object query embeddings (first stage).
        3. Predict the final classification labels and refined bounding boxes using the decoder's last
           hidden state (second stage).
        4. Pass the high-resolution spatial features and query hidden states through the segmentation
           head to produce the final, detailed instance masks.
    """
    )
    def forward(
        self,
        pixel_values: torch.FloatTensor = None,
        pixel_mask: torch.LongTensor | None = None,
        labels: list[dict] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> dict[str, torch.Tensor]:
        r"""
        labels (`list[Dict]` of len `(batch_size,)`, *optional*):
            Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
            following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
            respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
            in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
        """
        image_size = pixel_values.shape[-2:]

        # Step 1.
        outputs = self.model.model(pixel_values, pixel_mask=pixel_mask, **kwargs)

        spatial_features = outputs.backbone_features
        last_hidden_states = outputs.last_hidden_state
        intermediate_reference_points = outputs.intermediate_reference_points
        enc_outputs_class = outputs.enc_outputs_class

        # Step 2.
        enc_outputs_class_logits = self.model.predict_encoder_class_logits(enc_outputs_class)
        enc_outputs_masks = self.segmentation_head(spatial_features, enc_outputs_class, image_size, skip_blocks=True)

        # Step 3.
        logits, pred_boxes = self.model.predict_class_and_boxes(last_hidden_states, intermediate_reference_points[-1])

        # Step 4.
        outputs_masks = self.segmentation_head(spatial_features, outputs.intermediate_hidden_states, image_size)
        pred_masks = outputs_masks[-1]

        loss, loss_dict, auxiliary_outputs = None, None, None
        if labels is not None:
            outputs_class, outputs_coord = None, None
            if self.config.auxiliary_loss:
                outputs_class, outputs_coord = self.model.predict_class_and_boxes(
                    outputs.intermediate_hidden_states, intermediate_reference_points
                )
            loss, loss_dict, auxiliary_outputs = self.loss_function(
                logits,
                labels,
                self.device,
                pred_boxes,
                pred_masks,
                self.config,
                outputs_class,
                outputs_coord,
                outputs_masks,
                enc_outputs_class_logits,
                outputs.enc_outputs_coord_logits,
                enc_outputs_masks,
            )

        return RfDetrInstanceSegmentationOutput(
            loss=loss,
            loss_dict=loss_dict,
            logits=logits,
            pred_boxes=pred_boxes,
            pred_masks=pred_masks,
            auxiliary_outputs=auxiliary_outputs,
            last_hidden_state=outputs.last_hidden_state,
            intermediate_hidden_states=outputs.intermediate_hidden_states,
            intermediate_reference_points=outputs.intermediate_reference_points,
            init_reference_points=outputs.init_reference_points,
            enc_outputs_mask_logits=enc_outputs_masks,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )


__all__ = [
    "RfDetrConfig",
    "RfDetrModel",
    "RfDetrForObjectDetection",
    "RfDetrForInstanceSegmentation",
    "RfDetrPreTrainedModel",
    "RfDetrDinov2Config",
    "RfDetrDinov2Backbone",
    "RfDetrDinov2PreTrainedModel",
    "RfDetrImageProcessor",
]
