onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/ops/image.py

@@ -0,0 +1,748 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Image and spatial transformation operators.
+
+This module implements ONNX operators for image resizing and
+spatial dimension rearrangement.
+"""
+
+from typing import TYPE_CHECKING
+
+import onnx
+import torch
+
+from ..op_registry import register
+from ..utils.attributes import get_attribute
+from ..utils.op_helpers import get_optional_input
+
+if TYPE_CHECKING:
+    from ..graph_builder import GraphBuilder
+
+
+# =============================================================================
+# Resize operator
+# =============================================================================
+
+
+@register("Resize")
+def resize(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Resize tensor using interpolation."""
+    x = builder.get_value(node.input[0])
+
+    # roi, scales, sizes are optional inputs
+    roi = get_optional_input(builder, node, 1)
+    scales = get_optional_input(builder, node, 2)
+    sizes = get_optional_input(builder, node, 3)
+
+    mode = get_attribute(node, "mode", "nearest")
+    coordinate_transformation_mode = get_attribute(
+        node, "coordinate_transformation_mode", "half_pixel"
+    )
+
+    def _resize(x, roi, scales, sizes, mode, coord_mode):
+        import torch.nn.functional as F
+
+        # Map ONNX mode to PyTorch mode
+        mode_map = {
+            "nearest": "nearest",
+            "linear": "bilinear" if x.dim() == 4 else "linear",
+            "cubic": "bicubic",
+        }
+        torch_mode = mode_map.get(mode, "nearest")
+
+        # Determine align_corners based on coordinate transformation mode
+        align_corners = coord_mode == "align_corners"
+        if torch_mode == "nearest":
+            align_corners = None
+
+        if sizes is not None:
+            # Use explicit sizes
+            size_list = sizes.tolist() if isinstance(sizes, torch.Tensor) else sizes
+            # Skip batch and channel dimensions
+            output_size = [int(s) for s in size_list[2:]]
+        elif scales is not None:
+            # Use scales
+            scale_list = scales.tolist() if isinstance(scales, torch.Tensor) else scales
+            input_shape = x.shape[2:]
+            output_size = [int(s * sc) for s, sc in zip(input_shape, scale_list[2:])]
+        else:
+            return x
+
+        kwargs = {"size": output_size, "mode": torch_mode}
+        if align_corners is not None and torch_mode not in ["nearest", "area"]:
+            kwargs["align_corners"] = align_corners
+
+        return F.interpolate(x, **kwargs)
+
+    return builder.call_function(
+        _resize, args=(x, roi, scales, sizes, mode, coordinate_transformation_mode)
+    )
+
+
+# =============================================================================
+# Upsample operator (deprecated in opset 10, replaced by Resize)
+# =============================================================================
+
+
+@register("Upsample", since_version=7)
+def upsample(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Upsample tensor using interpolation.
+
+    Deprecated: This operator is deprecated since opset 10.
+    Use Resize operator instead.
+
+    Opset 7-8: scales is an attribute
+    Opset 9: scales is an input
+    """
+    x = builder.get_value(node.input[0])
+
+    # In opset 9, scales is an input; in opset 7-8, it's an attribute
+    opset = builder.opset_version
+    if opset >= 9 and len(node.input) > 1 and node.input[1]:
+        scales = builder.get_value(node.input[1])
+    else:
+        scales = get_attribute(node, "scales")
+
+    mode = get_attribute(node, "mode", "nearest")
+
+    def _upsample(x, scales, mode):
+        import torch.nn.functional as F
+
+        # Map ONNX mode to PyTorch mode
+        mode_map = {
+            "nearest": "nearest",
+            "linear": "bilinear" if x.dim() == 4 else "linear",
+            "cubic": "bicubic",
+        }
+        torch_mode = mode_map.get(mode, "nearest")
+
+        # Use scales to compute output size
+        scale_list = scales.tolist() if isinstance(scales, torch.Tensor) else scales
+        input_shape = x.shape[2:]
+        output_size = [int(s * sc) for s, sc in zip(input_shape, scale_list[2:])]
+
+        kwargs = {"size": output_size, "mode": torch_mode}
+        # align_corners is not used for nearest mode
+        if torch_mode not in ["nearest", "area"]:
+            kwargs["align_corners"] = False
+
+        return F.interpolate(x, **kwargs)
+
+    return builder.call_function(_upsample, args=(x, scales, mode))
+
+
+# =============================================================================
+# Depth/Space rearrangement operators
+# =============================================================================
+
+
+@register("DepthToSpace")
+def depth_to_space(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Rearrange depth to spatial dimensions."""
+    x = builder.get_value(node.input[0])
+    blocksize = get_attribute(node, "blocksize")
+    mode = get_attribute(node, "mode", "DCR")
+
+    def _depth_to_space(x, blocksize, mode):
+        b, c, h, w = x.shape
+        if mode == "DCR":
+            # Depth-Column-Row
+            x = x.reshape(b, blocksize, blocksize, c // (blocksize**2), h, w)
+            x = x.permute(0, 3, 4, 1, 5, 2)
+            x = x.reshape(b, c // (blocksize**2), h * blocksize, w * blocksize)
+        else:
+            # CRD mode
+            x = x.reshape(b, c // (blocksize**2), blocksize, blocksize, h, w)
+            x = x.permute(0, 1, 4, 2, 5, 3)
+            x = x.reshape(b, c // (blocksize**2), h * blocksize, w * blocksize)
+        return x
+
+    return builder.call_function(_depth_to_space, args=(x, blocksize, mode))
+
+
+@register("SpaceToDepth")
+def space_to_depth(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Rearrange spatial dimensions to depth."""
+    x = builder.get_value(node.input[0])
+    blocksize = get_attribute(node, "blocksize")
+
+    def _space_to_depth(x, blocksize):
+        b, c, h, w = x.shape
+        x = x.reshape(b, c, h // blocksize, blocksize, w // blocksize, blocksize)
+        x = x.permute(0, 3, 5, 1, 2, 4)
+        x = x.reshape(b, c * blocksize * blocksize, h // blocksize, w // blocksize)
+        return x
+
+    return builder.call_function(_space_to_depth, args=(x, blocksize))
+
+
+# =============================================================================
+# Col2Im operator
+# =============================================================================
+
+
+@register("Col2Im", since_version=18)
+def col2im(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Rearrange column blocks back into a multidimensional image.
+
+    ONNX Col2Im is the inverse of Im2Col. It combines sliding local blocks
+    (columns) back into a larger image tensor.
+
+    Inputs:
+        input: [N, C * prod(block_shape), L] - batched column data
+        image_shape: spatial dimensions of the output image
+        block_shape: shape of the sliding block
+
+    Attributes:
+        strides: stride along each spatial axis (default: 1)
+        pads: padding for each spatial axis in ONNX format (default: 0)
+        dilations: dilation for each spatial axis (default: 1)
+
+    Output:
+        [N, C, *image_shape] - reconstructed image
+    """
+    x = builder.get_value(node.input[0])
+    image_shape = builder.get_value(node.input[1])
+    block_shape = builder.get_value(node.input[2])
+
+    strides = get_attribute(node, "strides")
+    pads = get_attribute(node, "pads")
+    dilations = get_attribute(node, "dilations")
+
+    def _col2im(x, image_shape, block_shape, strides, pads, dilations):
+        import torch.nn.functional as F
+        from functools import reduce
+        from itertools import product
+        from operator import mul
+
+        # Convert to lists if tensors
+        if isinstance(image_shape, torch.Tensor):
+            image_shape = image_shape.tolist()
+        if isinstance(block_shape, torch.Tensor):
+            block_shape = block_shape.tolist()
+
+        n_dims = len(block_shape)
+
+        # Default values
+        if strides is None:
+            strides = [1] * n_dims
+        if pads is None:
+            pads = [0] * (2 * n_dims)
+        if dilations is None:
+            dilations = [1] * n_dims
+
+        # For 2D, use PyTorch's optimized fold
+        if n_dims == 2:
+            # PyTorch fold uses symmetric padding per dimension
+            padding = (pads[0], pads[1])
+            return F.fold(
+                x,
+                output_size=tuple(image_shape),
+                kernel_size=tuple(block_shape),
+                stride=tuple(strides),
+                padding=padding,
+                dilation=tuple(dilations),
+            )
+
+        # For N-D, implement manually
+        N = x.shape[0]
+        L = x.shape[2]
+        block_size = reduce(mul, block_shape, 1)
+        C = x.shape[1] // block_size
+
+        # Reshape input: [N, C * prod(block_shape), L] -> [N, C, *block_shape, L]
+        input_reshaped = x.reshape(N, C, *block_shape, L)
+
+        # Initialize output: [N, C, *image_shape]
+        output = torch.zeros(N, C, *image_shape, dtype=x.dtype, device=x.device)
+
+        # Compute effective kernel size after dilation
+        effective_block = [(b - 1) * d + 1 for b, d in zip(block_shape, dilations)]
+
+        # Compute number of blocks in each dimension
+        n_blocks = []
+        for i, (img_dim, eff_block, s) in enumerate(
+            zip(image_shape, effective_block, strides)
+        ):
+            p_begin = pads[i]
+            p_end = pads[n_dims + i]
+            n_block = (img_dim + p_begin + p_end - eff_block) // s + 1
+            n_blocks.append(n_block)
+
+        # Iterate over all block positions
+        block_indices = list(product(*[range(nb) for nb in n_blocks]))
+
+        for l_idx, block_idx in enumerate(block_indices):
+            # Compute starting position for this block
+            starts = [
+                bi * s - pads[i] for i, (bi, s) in enumerate(zip(block_idx, strides))
+            ]
+
+            # For each position in the block
+            block_positions = list(product(*[range(b) for b in block_shape]))
+            for block_pos in block_positions:
+                # Compute actual output position with dilation
+                output_pos = [
+                    starts[i] + block_pos[i] * dilations[i] for i in range(n_dims)
+                ]
+
+                # Check bounds
+                valid = all(0 <= output_pos[i] < image_shape[i] for i in range(n_dims))
+                if valid:
+                    # Get value from input_reshaped: [N, C, *block_shape, L]
+                    idx = (slice(None), slice(None)) + tuple(block_pos) + (l_idx,)
+                    value = input_reshaped[idx]
+
+                    # Add to output
+                    out_idx = (slice(None), slice(None)) + tuple(output_pos)
+                    output[out_idx] += value
+
+        return output
+
+    return builder.call_function(
+        _col2im, args=(x, image_shape, block_shape, strides, pads, dilations)
+    )
+
+
+# =============================================================================
+# CenterCropPad operator
+# =============================================================================
+
+
+# =============================================================================
+# GridSample operator
+# =============================================================================
+
+
+@register("GridSample", since_version=16)
+def grid_sample(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Sample input using grid of sampling locations.
+
+    Given an input X and a flow-field grid, computes the output Y using X values
+    and pixel locations from the grid. This is equivalent to PyTorch's
+    torch.nn.functional.grid_sample.
+
+    Inputs:
+        X: Input tensor of shape (N, C, H, W) for 4D or (N, C, D, H, W) for 5D
+        grid: Grid tensor of shape (N, H_out, W_out, 2) for 4D or
+              (N, D_out, H_out, W_out, 3) for 5D
+
+    Attributes:
+        align_corners: If 1, extrema (-1 and 1) refer to center of corner pixels.
+                      If 0, they refer to corner points. Default: 0
+        mode: Interpolation mode - 'linear'/'bilinear' (default), 'nearest',
+              'cubic'/'bicubic'. Opset 16 uses bilinear/bicubic, opset 20+ uses
+              linear/cubic.
+        padding_mode: Padding mode for outside grid values - 'zeros' (default),
+                     'border', 'reflection'
+
+    Output:
+        Y: Output tensor of shape (N, C, H_out, W_out) or (N, C, D_out, H_out, W_out)
+    """
+    x = builder.get_value(node.input[0])
+    grid = builder.get_value(node.input[1])
+
+    align_corners = get_attribute(node, "align_corners", 0)
+    # Handle different mode names across opset versions
+    # Opset 16: bilinear (default), nearest, bicubic
+    # Opset 20+: linear (default), nearest, cubic
+    mode = get_attribute(node, "mode", "linear")
+    padding_mode = get_attribute(node, "padding_mode", "zeros")
+
+    def _grid_sample(x, grid, mode, padding_mode, align_corners):
+        import torch.nn.functional as F
+
+        # Map ONNX mode names to PyTorch mode names
+        # PyTorch expects: 'bilinear', 'nearest', 'bicubic' for 4D input
+        # PyTorch expects: 'bilinear', 'nearest' for 5D input (no bicubic)
+        mode_map = {
+            "linear": "bilinear",
+            "bilinear": "bilinear",
+            "nearest": "nearest",
+            "cubic": "bicubic",
+            "bicubic": "bicubic",
+        }
+        torch_mode = mode_map.get(mode, "bilinear")
+
+        # Convert align_corners from int to bool
+        align_corners_bool = bool(align_corners)
+
+        return F.grid_sample(
+            x,
+            grid,
+            mode=torch_mode,
+            padding_mode=padding_mode,
+            align_corners=align_corners_bool,
+        )
+
+    return builder.call_function(
+        _grid_sample, args=(x, grid, mode, padding_mode, align_corners)
+    )
+
+
+# =============================================================================
+# AffineGrid operator
+# =============================================================================
+
+
+@register("AffineGrid", since_version=20)
+def affine_grid(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Generate 2D or 3D flow field (sampling grid) from affine matrices.
+
+    Given a batch of affine matrices theta, generates a grid of sampling
+    locations. This is typically used with GridSample to build Spatial
+    Transformer Networks.
+
+    Inputs:
+        theta: Input batch of affine matrices with shape (N, 2, 3) for 2D
+               or (N, 3, 4) for 3D
+        size: Target output image size (N, C, H, W) for 2D or (N, C, D, H, W)
+              for 3D, as a 1-D tensor
+
+    Attributes:
+        align_corners: If 1, consider -1 and 1 to refer to the centers of the
+                      corner pixels. If 0, consider -1 and 1 to refer to the
+                      outer edge of corner pixels. Default: 0
+
+    Output:
+        grid: Output tensor of shape (N, H, W, 2) for 2D sample coordinates
+              or (N, D, H, W, 3) for 3D sample coordinates
+    """
+    theta = builder.get_value(node.input[0])
+    size = builder.get_value(node.input[1])
+
+    align_corners = get_attribute(node, "align_corners", 0)
+
+    def _affine_grid(theta, size, align_corners):
+        import torch.nn.functional as F
+
+        # Convert size tensor to a list of integers for torch.Size
+        if isinstance(size, torch.Tensor):
+            size_list = size.tolist()
+        else:
+            size_list = list(size)
+        size_tuple = torch.Size([int(s) for s in size_list])
+
+        # Convert align_corners from int to bool
+        align_corners_bool = bool(align_corners)
+
+        return F.affine_grid(theta, size_tuple, align_corners=align_corners_bool)
+
+    return builder.call_function(_affine_grid, args=(theta, size, align_corners))
+
+
+# =============================================================================
+# CenterCropPad operator
+# =============================================================================
+
+
+@register("CenterCropPad", since_version=18)
+def center_crop_pad(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Center crop or pad an input to given dimensions.
+
+    The crop/pad dimensions can be specified for a subset of the axes.
+    Unspecified dimensions will remain unchanged.
+
+    For cropping (input > target): centered window, start position rounded down.
+    For padding (input < target): centered padding, extra pixel on right if odd.
+
+    Inputs:
+        input_data: Input tensor to crop/pad
+        shape: 1-D tensor of target dimensions for specified axes
+
+    Attributes:
+        axes: Subset of axes that shape refers to (default: all axes)
+
+    Output:
+        Output tensor with specified dimensions
+    """
+    x = builder.get_value(node.input[0])
+    shape = builder.get_value(node.input[1])
+
+    axes = get_attribute(node, "axes")
+
+    def _center_crop_pad(x, shape, axes):
+        # Convert shape to list if tensor
+        if isinstance(shape, torch.Tensor):
+            target_shape = shape.tolist()
+        else:
+            target_shape = list(shape)
+
+        ndim = x.dim()
+
+        # If axes is not provided, use all axes
+        if axes is None:
+            axes_list = list(range(ndim))
+        else:
+            axes_list = list(axes)
+
+        # Normalize negative axes
+        axes_list = [(a + ndim) if a < 0 else a for a in axes_list]
+
+        # Build slices for cropping and padding amounts
+        result = x
+        for i, axis in enumerate(axes_list):
+            current_size = result.shape[axis]
+            target_size = int(target_shape[i])
+
+            if current_size == target_size:
+                # No change needed for this axis
+                continue
+            elif current_size > target_size:
+                # Crop: extract centered window
+                # Start position is rounded down (floor division)
+                diff = current_size - target_size
+                start = diff // 2
+                end = start + target_size
+
+                # Build slice for this axis
+                slices = [slice(None)] * result.dim()
+                slices[axis] = slice(start, end)
+                result = result[tuple(slices)]
+            else:
+                # Pad: add zeros centered
+                # Extra pixel goes to the right side
+                diff = target_size - current_size
+                pad_before = diff // 2
+                pad_after = diff - pad_before
+
+                # torch.nn.functional.pad uses reverse order: last dim first
+                # and pairs are (before, after) for each dim from last to first
+                # We need to construct padding for just this one axis
+
+                # Number of dimensions from the end
+                dims_from_end = result.dim() - 1 - axis
+
+                # Build pad tuple: pairs for each dim from last to first
+                # We only pad the current axis
+                pad = [0] * (2 * result.dim())
+                # Index in pad list: dims_from_end * 2 for before, +1 for after
+                pad[dims_from_end * 2] = pad_before
+                pad[dims_from_end * 2 + 1] = pad_after
+
+                result = torch.nn.functional.pad(result, pad, mode="constant", value=0)
+
+        return result
+
+    return builder.call_function(_center_crop_pad, args=(x, shape, axes))
+
+
+# =============================================================================
+# RoiAlign operator
+# =============================================================================
+
+
+@register("RoiAlign")
+def roi_align(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Region of Interest (RoI) Align operator.
+
+    Performs RoI pooling with bilinear interpolation for sub-pixel accuracy.
+    """
+    x = builder.get_value(node.input[0])
+    rois = builder.get_value(node.input[1])
+    batch_indices = builder.get_value(node.input[2])
+
+    # Get attributes with defaults
+    mode = get_attribute(node, "mode", "avg")
+    output_height = get_attribute(node, "output_height", 1)
+    output_width = get_attribute(node, "output_width", 1)
+    sampling_ratio = get_attribute(node, "sampling_ratio", 0)
+    spatial_scale = get_attribute(node, "spatial_scale", 1.0)
+    coordinate_transformation_mode = get_attribute(
+        node, "coordinate_transformation_mode", "half_pixel"
+    )
+
+    # ONNX coordinate_transformation_mode:
+    # - "half_pixel": pixel shift by -0.5, corresponds to aligned=True in PyTorch
+    # - "output_half_pixel": no pixel shift (legacy), corresponds to aligned=False
+    aligned = coordinate_transformation_mode == "half_pixel"
+
+    def _roi_align(
+        x,
+        rois,
+        batch_indices,
+        mode,
+        output_height,
+        output_width,
+        sampling_ratio,
+        spatial_scale,
+        aligned,
+    ):
+        from torchvision.ops import roi_align as tv_roi_align
+
+        # PyTorch expects boxes in format [batch_idx, x1, y1, x2, y2]
+        boxes = torch.cat([batch_indices.unsqueeze(1).float(), rois.float()], dim=1)
+        output_size = (output_height, output_width)
+
+        if mode == "avg":
+            # Use torchvision's roi_align directly for average mode
+            # sampling_ratio: ONNX uses 0 for adaptive, PyTorch uses -1
+            torch_sampling = sampling_ratio if sampling_ratio > 0 else -1
+            return tv_roi_align(
+                x,
+                boxes,
+                output_size,
+                spatial_scale=spatial_scale,
+                sampling_ratio=torch_sampling,
+                aligned=aligned,
+            )
+        else:
+            # Max mode: ONNX defines max pooling differently from standard
+            # bilinear interpolation. For each sample point, it takes the
+            # max of the 4 weighted corner values (not the sum).
+            return _roi_align_max_mode(
+                x,
+                rois,
+                batch_indices,
+                output_height,
+                output_width,
+                sampling_ratio,
+                spatial_scale,
+                aligned,
+            )
+
+    return builder.call_function(
+        _roi_align,
+        args=(
+            x,
+            rois,
+            batch_indices,
+            mode,
+            output_height,
+            output_width,
+            sampling_ratio,
+            spatial_scale,
+            aligned,
+        ),
+    )
+
+
+def _roi_align_max_mode(
+    x,
+    rois,
+    batch_indices,
+    output_height,
+    output_width,
+    sampling_ratio,
+    spatial_scale,
+    half_pixel,
+):
+    """ONNX RoiAlign with max pooling mode.
+
+    For each output bin, samples at grid points and takes the MAX of the
+    weighted corner values at each sampling point, then MAX across all
+    sampling points in the bin.
+    """
+    num_rois = rois.shape[0]
+    channels = x.shape[1]
+    height = x.shape[2]
+    width = x.shape[3]
+
+    output = torch.zeros(
+        num_rois, channels, output_height, output_width, dtype=x.dtype, device=x.device
+    )
+
+    for n in range(num_rois):
+        roi_batch_ind = int(batch_indices[n].item())
+        roi = rois[n]
+
+        # Apply spatial scale and offset
+        offset = 0.5 if half_pixel else 0.0
+        roi_start_w = float(roi[0]) * spatial_scale - offset
+        roi_start_h = float(roi[1]) * spatial_scale - offset
+        roi_end_w = float(roi[2]) * spatial_scale - offset
+        roi_end_h = float(roi[3]) * spatial_scale - offset
+
+        roi_width = roi_end_w - roi_start_w
+        roi_height = roi_end_h - roi_start_h
+
+        if not half_pixel:
+            # Force malformed ROIs to be 1x1
+            roi_width = max(roi_width, 1.0)
+            roi_height = max(roi_height, 1.0)
+
+        bin_size_h = roi_height / output_height
+        bin_size_w = roi_width / output_width
+
+        # Determine sampling grid size
+        if sampling_ratio > 0:
+            roi_bin_grid_h = sampling_ratio
+            roi_bin_grid_w = sampling_ratio
+        else:
+            roi_bin_grid_h = int(
+                torch.ceil(torch.tensor(roi_height / output_height)).item()
+            )
+            roi_bin_grid_w = int(
+                torch.ceil(torch.tensor(roi_width / output_width)).item()
+            )
+            roi_bin_grid_h = max(1, roi_bin_grid_h)
+            roi_bin_grid_w = max(1, roi_bin_grid_w)
+
+        for c in range(channels):
+            for ph in range(output_height):
+                for pw in range(output_width):
+                    output_val = None
+
+                    for iy in range(roi_bin_grid_h):
+                        yy = (
+                            roi_start_h
+                            + ph * bin_size_h
+                            + (iy + 0.5) * bin_size_h / roi_bin_grid_h
+                        )
+                        for ix in range(roi_bin_grid_w):
+                            xx = (
+                                roi_start_w
+                                + pw * bin_size_w
+                                + (ix + 0.5) * bin_size_w / roi_bin_grid_w
+                            )
+
+                            # Check bounds
+                            if yy < -1.0 or yy > height or xx < -1.0 or xx > width:
+                                continue
+
+                            y = max(yy, 0.0)
+                            xc = max(xx, 0.0)
+
+                            y_low = int(y)
+                            x_low = int(xc)
+
+                            if y_low >= height - 1:
+                                y_high = y_low = height - 1
+                                y = float(y_low)
+                            else:
+                                y_high = y_low + 1
+
+                            if x_low >= width - 1:
+                                x_high = x_low = width - 1
+                                xc = float(x_low)
+                            else:
+                                x_high = x_low + 1
+
+                            ly = y - y_low
+                            lx = xc - x_low
+                            hy = 1.0 - ly
+                            hx = 1.0 - lx
+
+                            # Weights
+                            w1 = hy * hx
+                            w2 = hy * lx
+                            w3 = ly * hx
+                            w4 = ly * lx
+
+                            # Get corner values
+                            v1 = x[roi_batch_ind, c, y_low, x_low].item()
+                            v2 = x[roi_batch_ind, c, y_low, x_high].item()
+                            v3 = x[roi_batch_ind, c, y_high, x_low].item()
+                            v4 = x[roi_batch_ind, c, y_high, x_high].item()
+
+                            # ONNX max mode: max of weighted corners
+                            val = max(w1 * v1, w2 * v2, w3 * v3, w4 * v4)
+
+                            if output_val is None:
+                                output_val = val
+                            else:
+                                output_val = max(output_val, val)
+
+                    if output_val is not None:
+                        output[n, c, ph, pw] = output_val
+
+    return output