108 lines
4.4 KiB
Python
108 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (C) 2018-2025 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Factory functions for all openvino ops."""
|
|
from functools import partial
|
|
from typing import Optional
|
|
|
|
from openvino import Node
|
|
from openvino.utils.node_factory import _get_node_factory
|
|
from openvino.utils.decorators import nameable_op
|
|
from openvino.utils.types import (
|
|
NodeInput,
|
|
as_nodes,
|
|
)
|
|
|
|
_get_node_factory_opset11 = partial(_get_node_factory, "opset11")
|
|
|
|
# -------------------------------------------- ops ------------------------------------------------
|
|
|
|
|
|
@nameable_op
|
|
def interpolate(
|
|
image: NodeInput,
|
|
scales_or_sizes: NodeInput,
|
|
mode: str,
|
|
shape_calculation_mode: str,
|
|
pads_begin: Optional[list[int]] = None,
|
|
pads_end: Optional[list[int]] = None,
|
|
coordinate_transformation_mode: str = "half_pixel",
|
|
nearest_mode: str = "round_prefer_floor",
|
|
antialias: bool = False,
|
|
cube_coeff: float = -0.75,
|
|
axes: Optional[NodeInput] = None,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""Perfors the interpolation of the input tensor.
|
|
|
|
:param image: The node providing input tensor with data for interpolation.
|
|
:param scales_or_sizes:
|
|
1D tensor providing information used to calculate the output shape
|
|
of the operation. It might contain floats (scales) or integers(sizes).
|
|
:param mode: Specifies type of interpolation. Possible values are: nearest, linear,
|
|
linear_onnx, cubic, bilinear_pillow, bicubic_pillow.
|
|
:param shape_calculation_mode:
|
|
Specifies how the scales_or_sizes input should be interpreted.
|
|
:param pads_begin: Specifies the number of pixels to add to the beginning of the image
|
|
being interpolated. Default is None.
|
|
:param pads_end: Specifies the number of pixels to add to the end of the image being
|
|
interpolated. Default is None.
|
|
:param coordinate_transformation_mode:
|
|
Specifies how to transform the coordinate in the resized tensor to the
|
|
coordinate in the original tensor. Default is "half_pixel".
|
|
:param nearest_mode: Specifies round mode when mode == nearest and is used only when
|
|
mode == nearest. Default is "round_prefer_floor".
|
|
:param antialias: Specifies whether to perform anti-aliasing. Default is False.
|
|
:param cube_coeff: Specifies the parameter a for cubic interpolation. Default is -0.75.
|
|
:param axes: 1D tensor specifying dimension indices where interpolation is applied.
|
|
The default is None.
|
|
:param name: Optional name for the output node. The default is None.
|
|
:return: Node representing the interpolation operation.
|
|
"""
|
|
attrs = {
|
|
"mode": mode,
|
|
"shape_calculation_mode": shape_calculation_mode,
|
|
"coordinate_transformation_mode": coordinate_transformation_mode,
|
|
"nearest_mode": nearest_mode,
|
|
"antialias": antialias,
|
|
"cube_coeff": cube_coeff,
|
|
}
|
|
|
|
attrs["pads_begin"] = [] if pads_begin is None else pads_begin
|
|
attrs["pads_end"] = [] if pads_end is None else pads_end
|
|
|
|
inputs = as_nodes(image, scales_or_sizes, name=name) if axes is None else as_nodes(image, scales_or_sizes, axes, name=name)
|
|
|
|
return _get_node_factory_opset11().create("Interpolate", inputs, attrs)
|
|
|
|
|
|
@nameable_op
|
|
def topk(
|
|
data: NodeInput,
|
|
k: NodeInput,
|
|
axis: int,
|
|
mode: str,
|
|
sort: str,
|
|
index_element_type: str = "i32",
|
|
stable: bool = False,
|
|
name: Optional[str] = None,
|
|
) -> Node:
|
|
"""Return a node which performs TopK.
|
|
|
|
:param data: Input data.
|
|
:param k: K.
|
|
:param axis: TopK Axis.
|
|
:param mode: Compute TopK largest ('max') or smallest ('min')
|
|
:param sort: Order of output elements (sort by: 'none', 'index' or 'value')
|
|
:param index_element_type: Type of output tensor with indices.
|
|
:param stable: Specifies whether the equivalent elements should maintain
|
|
their relative order from the input tensor during sorting.
|
|
:return: The new node which performs TopK
|
|
"""
|
|
return _get_node_factory_opset11().create(
|
|
"TopK",
|
|
as_nodes(data, k, name=name),
|
|
{"axis": axis, "mode": mode, "sort": sort, "index_element_type": index_element_type, "stable": stable},
|
|
)
|