#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/pp_ocrv5_server_det/modular_pp_ocrv5_server_det.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_pp_ocrv5_server_det.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The PaddlePaddle Team and 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 math

import numpy as np
import torch
import torchvision.transforms.v2.functional as tvF

from ...feature_extraction_utils import BatchFeature
from ...image_processing_backends import TorchvisionBackend
from ...image_transforms import group_images_by_shape, reorder_images
from ...image_utils import PILImageResampling, SizeDict
from ...processing_utils import ImagesKwargs
from ...utils import auto_docstring, is_cv2_available, requires_backends
from ...utils.generic import TensorType
from ...utils.import_utils import requires


if is_cv2_available():
    import cv2


class PPOCRV5ServerDetImageProcessorKwargs(ImagesKwargs, total=False):
    r"""
    limit_side_len (`int`, *optional*, defaults to `960`):
        Maximum or minimum side length.
    limit_type (`str`, *optional*, defaults to `max`):
        Resizing strategy: "max", "min", or "resize_long".
    max_side_limit (`int`, *optional* defaults to `4000`):
        Maximum allowed side length.
    """

    limit_side_len: int
    limit_type: str
    max_side_limit: int


@auto_docstring
@requires(backends=("torch",))
class PPOCRV5ServerDetImageProcessor(TorchvisionBackend):
    resample = 2
    image_mean = [0.406, 0.456, 0.485]
    image_std = [0.225, 0.224, 0.229]
    size = {"height": 960, "width": 960}
    do_resize = True
    do_rescale = True
    do_normalize = True
    limit_side_len = 960
    limit_type = "max"
    max_side_limit = 4000
    valid_kwargs = PPOCRV5ServerDetImageProcessorKwargs

    def _preprocess(
        self,
        images: list["torch.Tensor"],
        do_resize: bool,
        size: SizeDict,
        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        limit_side_len: int,
        limit_type: str,
        max_side_limit: int,
        disable_grouping: bool | None,
        return_tensors: str | TensorType | None,
        **kwargs,
    ) -> BatchFeature:
        target_sizes = []

        # Group images by their original spatial shape to enable batched resizing (optimization for efficiency)
        # [Key Change] Unlike the original implementation, we now track target shapes for each original shape group
        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
        # Store resized image batches mapped to their original shape keys
        resized_images_grouped = {}
        # [Key Change] Core addition: Mapping from original image shape to target resize shape
        # This dict ensures consistent target shape handling across all subsequent operations (resize/processing)
        target_shape_per_shape = {}
        for shape, stacked_images in grouped_images.items():
            if do_resize:
                resize_size, target_shape = self.get_image_size(
                    stacked_images[0], limit_side_len, limit_type, max_side_limit
                )
                target_shape_per_shape[shape] = target_shape
                stacked_images = self.resize(image=stacked_images.float(), size=resize_size, resample=resample)
            resized_images_grouped[shape] = stacked_images

        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
        if do_resize:
            target_sizes = [target_shape_per_shape[grouped_images_index[i][0]] for i in range(len(images))]

        # Group images by size for further processing
        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
        processed_images_grouped = {}
        for shape, stacked_images in grouped_images.items():
            stacked_images = self.rescale_and_normalize(
                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
            )
            # BGR to RGB conversion
            stacked_images = stacked_images[:, [2, 1, 0], :, :]
            processed_images_grouped[shape] = stacked_images

        pixel_values = reorder_images(processed_images_grouped, grouped_images_index)

        return BatchFeature(
            data={"pixel_values": pixel_values, "target_sizes": target_sizes},
            tensor_type=return_tensors,
        )

    def _unclip(self, contour_box, unclip_ratio):
        """
        Expands (dilates) a detected text bounding box to recover the full text region.

        Args:
            contour_box (np.ndarray): Input contour of shape (N, 2), where N is the number of points.
            unclip_ratio (float): Expansion ratio, typically greater than 1.0.

        Returns:
            np.ndarray: Expanded contour of shape (M, 2).
        """
        # --- 1. Parameter calculation ---
        polygon = contour_box.reshape(-1, 2).astype(np.float32)
        perimeter = cv2.arcLength(polygon, True)
        area = cv2.contourArea(polygon)
        offset_distance = area * unclip_ratio / perimeter

        # --- 2. Determine polygon orientation and edge normals ---
        x, y = polygon[:, 0], polygon[:, 1]
        is_counter_clockwise = (x @ np.roll(y, -1) - y @ np.roll(x, -1)) > 0.0

        edges = np.roll(polygon, -1, axis=0) - polygon
        edge_lengths = np.linalg.norm(edges, axis=1, keepdims=True)
        edge_directions = edges / np.maximum(edge_lengths, 1e-6)

        if is_counter_clockwise:
            normals = np.stack([edge_directions[:, 1], -edge_directions[:, 0]], axis=1)
        else:
            normals = np.stack([-edge_directions[:, 1], edge_directions[:, 0]], axis=1)

        # --- 3. Calculate new vertices from intersecting shifted edge lines ---
        shifted_points = polygon + offset_distance * normals

        prev_shifted_points = np.roll(shifted_points, 1, axis=0)
        prev_edge_directions = np.roll(edge_directions, 1, axis=0)

        cross_product = (
            prev_edge_directions[:, 0] * edge_directions[:, 1] - prev_edge_directions[:, 1] * edge_directions[:, 0]
        )

        is_parallel_mask = np.abs(cross_product) < 1e-6
        cross_product_safe = np.where(is_parallel_mask, 1.0, cross_product)

        vec_to_current = shifted_points - prev_shifted_points
        intersection_param = (
            vec_to_current[:, 0] * edge_directions[:, 1] - vec_to_current[:, 1] * edge_directions[:, 0]
        ) / cross_product_safe

        new_vertices = prev_shifted_points + prev_edge_directions * intersection_param[:, None]

        # --- 4. Handle near-parallel adjacent edges with a fallback ---
        if np.any(is_parallel_mask):
            prev_normals = np.roll(normals, 1, axis=0)
            fallback_points = polygon + 0.5 * offset_distance * (prev_normals + normals)
            new_vertices[is_parallel_mask] = fallback_points[is_parallel_mask]

        return np.array([new_vertices.astype(np.float32)])

    def _get_mini_boxes(self, contour):
        """
        Computes the minimum-area bounding rectangle for a given contour and returns
        its four corners in a consistent order (top-left, bottom-left, bottom-right, top-right).

        Args:
            contour (np.ndarray): Input contour of shape (N, 1, 2).

        Returns:
            tuple:
                - box (list): List of four corner points in order.
                - short_side_length (float): Length of the shorter side of the bounding rectangle.
        """
        bounding_box = cv2.minAreaRect(contour)
        points = sorted(cv2.boxPoints(bounding_box), key=lambda x: x[0])

        index_1, index_2, index_3, index_4 = 0, 1, 2, 3
        if points[1][1] > points[0][1]:
            index_1 = 0
            index_4 = 1
        else:
            index_1 = 1
            index_4 = 0
        if points[3][1] > points[2][1]:
            index_2 = 2
            index_3 = 3
        else:
            index_2 = 3
            index_3 = 2

        box = [points[index_1], points[index_2], points[index_3], points[index_4]]
        return box, min(bounding_box[1])

    def _get_box_score(self, bitmap: np.ndarray, polygon_bounding_box: np.ndarray):
        """
        Computes the mean score of a bounding box region in the prediction map using
        a fast approach with axis-aligned bounding boxes.

        Args:
            bitmap (np.ndarray): Binary or float prediction map of shape (H, W).
            polygon_bounding_box (np.ndarray): Bounding box polygon of shape (N, 2).

        Returns:
            float: Mean score within the bounding box region.
        """
        height, width = bitmap.shape[:2]
        box = polygon_bounding_box.copy()
        xmin = max(0, min(math.floor(box[:, 0].min()), width - 1))
        xmax = max(0, min(math.ceil(box[:, 0].max()), width - 1))
        ymin = max(0, min(math.floor(box[:, 1].min()), height - 1))
        ymax = max(0, min(math.ceil(box[:, 1].max()), height - 1))

        mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
        box[:, 0] = box[:, 0] - xmin
        box[:, 1] = box[:, 1] - ymin
        cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
        return cv2.mean(bitmap[ymin : ymax + 1, xmin : xmax + 1], mask)[0]

    def _boxes_from_bitmap(
        self,
        prediction: np.ndarray,
        bitmap: np.ndarray,
        dest_width: int,
        dest_height: int,
        box_threshold: float,
        unclip_ratio: float,
        min_size: int,
        max_candidates: int,
    ):
        """
        Extracts axis-aligned or rotated bounding boxes from a binary segmentation map.

        Args:
            prediction (np.ndarray): Raw prediction map of shape (H, W).
            bitmap (np.ndarray): Binarized segmentation map of shape (H, W).
            dest_width (int): Original image width for scaling back.
            dest_height (int): Original image height for scaling back.
            box_threshold (float): Score threshold for filtering low-confidence boxes.
            unclip_ratio (float): Expansion ratio for contour unclipping.
            min_size (int): Minimum side length of valid boxes.
            max_candidates (int): Maximum number of contours to process.

        Returns:
            tuple:
                - boxes (np.ndarray): Array of boxes, each of shape (4, 2).
                - scores (list): List of corresponding scores.
        """

        height, width = bitmap.shape
        width_scale = dest_width / width
        height_scale = dest_height / height

        outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        contours = outs[1] if len(outs) == 3 else outs[0]

        num_contours = min(len(contours), max_candidates)

        boxes = []
        scores = []
        for index in range(num_contours):
            contour = contours[index]
            points, short_side_length = self._get_mini_boxes(contour)
            if short_side_length < min_size:
                continue
            points = np.array(points)
            score = self._get_box_score(prediction, points.reshape(-1, 2))
            if box_threshold > score:
                continue
            box = self._unclip(points, unclip_ratio).reshape(-1, 1, 2)
            box, short_side_length = self._get_mini_boxes(box)
            if short_side_length < min_size + 2:
                continue

            box = np.array(box)
            for i in range(box.shape[0]):
                box[i, 0] = max(0, min(round(box[i, 0] * width_scale), dest_width))
                box[i, 1] = max(0, min(round(box[i, 1] * height_scale), dest_height))

            boxes.append(box.astype(np.int16))
            scores.append(score)
        return np.array(boxes, dtype=np.int16), scores

    def get_image_size(
        self,
        image: "torch.Tensor",
        limit_side_len: int,
        limit_type: str,
        max_side_limit: int,
    ):
        """
        Computes the target size for resizing an image while preserving aspect ratio.

        Args:
            image (torch.Tensor): Input image.
            limit_side_len (int): Maximum or minimum side length.
            limit_type (str): Resizing strategy: "max", "min", or "resize_long".
            max_side_limit (int): Maximum allowed side length.

        Returns:
            tuple:
                - SizeDict: Target size.
                - torch.Tensor: Original size.
        """
        _, height, width = image.shape
        height, width = int(height), int(width)

        if limit_type == "max":
            if max(height, width) > limit_side_len:
                ratio = float(limit_side_len) / max(height, width)
            else:
                ratio = 1.0
        elif limit_type == "min":
            if min(height, width) < limit_side_len:
                ratio = float(limit_side_len) / min(height, width)
            else:
                ratio = 1.0
        elif limit_type == "resize_long":
            ratio = float(limit_side_len) / max(height, width)
        else:
            raise ValueError(f"PPOCRV5ServerDetImageProcessor does not support limit type: {limit_type}")

        resize_height = int(height * ratio)
        resize_width = int(width * ratio)

        if max_side_limit is not None and max(resize_height, resize_width) > max_side_limit:
            ratio = float(max_side_limit) / max(resize_height, resize_width)
            resize_height = int(resize_height * ratio)
            resize_width = int(resize_width * ratio)

        resize_height = max(int(round(resize_height / 32) * 32), 32)
        resize_width = max(int(round(resize_width / 32) * 32), 32)

        if resize_height == height and resize_width == width:
            return SizeDict(height=resize_height, width=resize_width), torch.tensor(
                [height, width], dtype=torch.float32, device=image.device
            )

        if resize_width <= 0 or resize_height <= 0:
            return None, (None, None)

        return SizeDict(height=resize_height, width=resize_width), torch.tensor(
            [height, width], dtype=torch.float32, device=image.device
        )

    def post_process_object_detection(
        self,
        predictions,
        threshold: float = 0.3,
        target_sizes: list[tuple[int, int]] | TensorType | None = None,
        box_threshold: float = 0.6,
        max_candidates: int = 1000,
        min_size: int = 3,
        unclip_ratio: float = 1.5,
    ):
        """
        Converts model outputs into detected text boxes in corners format (xmin, ymin, xmax, ymax).

        Args:
            predictions: Model outputs with `logits` attribute (probability maps of shape `(batch_size, 1, H, W)`).
            threshold (float): Binarization threshold.
            target_sizes: Original image sizes (height, width) per image.
            box_threshold (float): Box score threshold.
            max_candidates (int): Maximum number of boxes.
            min_size (int): Minimum box size.
            unclip_ratio (float): Expansion ratio.

        Returns:
            list[dict]: List of detection results per image. Each dict contains:
                - "boxes": `torch.Tensor` of shape `(N, 4)` in corners format (xmin, ymin, xmax, ymax)
                - "scores": `torch.Tensor` of shape `(N,)`
                - "labels": `torch.Tensor` of shape `(N,)` (class id 0 for text)
        """
        requires_backends(self, ["torch", "cv2"])
        if target_sizes is None:
            raise ValueError("target_sizes must be provided for post_process_object_detection")

        device = predictions.last_hidden_state.device
        results = []
        for prediction, size in zip(predictions.last_hidden_state, target_sizes):
            prediction = prediction[0, :, :].cpu().detach().numpy()
            size = size.cpu().detach().numpy()
            src_height, src_width = size
            mask = prediction > threshold
            boxes, scores = self._boxes_from_bitmap(
                prediction, mask, src_width, src_height, box_threshold, unclip_ratio, min_size, max_candidates
            )

            results.append(
                {
                    "boxes": torch.from_numpy(boxes).to(device),
                    "scores": torch.tensor(scores, dtype=torch.float32, device=device),
                    "labels": torch.zeros(len(scores), dtype=torch.long, device=device),  # Single class: text
                }
            )
        return results


__all__ = ["PPOCRV5ServerDetImageProcessor"]
