onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/ops/signal.py

@@ -0,0 +1,444 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Signal processing operators.
+
+This module implements specialized ONNX operators for signal processing,
+including STFT, mel spectrograms, window functions, and non-maximum suppression.
+"""
+
+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
+
+
+# =============================================================================
+# STFT (Short-time Fourier Transform)
+# =============================================================================
+
+
+@register("STFT", since_version=17)
+def stft(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Short-time Fourier Transform.
+
+    ONNX STFT operator computes the STFT of the input signal.
+
+    Inputs:
+        signal: [batch_size, signal_length, 1] for real or [batch_size, signal_length, 2] for complex
+        frame_step: scalar, the hop length
+        window (optional): 1D window tensor
+        frame_length (optional): scalar, the FFT size
+
+    Attributes:
+        onesided: int (default 1), whether to return one-sided output
+
+    Output:
+        [batch_size, frames, dft_unique_bins, 2] with real and imaginary components
+    """
+    signal = builder.get_value(node.input[0])
+    frame_step = builder.get_value(node.input[1])
+
+    # Optional window input
+    window = get_optional_input(builder, node, 2)
+
+    # Optional frame_length input
+    frame_length = get_optional_input(builder, node, 3)
+
+    # Get onesided attribute (default is 1)
+    onesided = get_attribute(node, "onesided", 1)
+
+    def _stft(signal, frame_step, window, frame_length, onesided):
+        # ONNX signal shape: [batch, signal_length, 1] (real) or [batch, signal_length, 2] (complex)
+        # We need to convert to PyTorch format: [batch, signal_length]
+
+        # Check if input is complex (last dim is 2)
+        is_complex_input = signal.shape[-1] == 2
+
+        if is_complex_input:
+            # Convert to complex tensor
+            signal_2d = torch.complex(signal[..., 0], signal[..., 1])
+        else:
+            # Squeeze the last dimension for real input
+            signal_2d = signal.squeeze(-1)
+
+        # Get scalar values
+        hop_length = (
+            int(frame_step.item())
+            if isinstance(frame_step, torch.Tensor)
+            else int(frame_step)
+        )
+
+        # Determine n_fft
+        if frame_length is not None:
+            n_fft = (
+                int(frame_length.item())
+                if isinstance(frame_length, torch.Tensor)
+                else int(frame_length)
+            )
+        elif window is not None:
+            n_fft = window.shape[0]
+        else:
+            raise ValueError("Either frame_length or window must be provided for STFT")
+
+        # Determine onesided behavior
+        # For complex input, onesided must be False
+        onesided_bool = bool(onesided) and not is_complex_input
+
+        # Call PyTorch stft
+        # PyTorch stft returns [batch, n_fft, frames] (complex) or [batch, n_fft, frames, 2] (real)
+        result = torch.stft(
+            signal_2d,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=n_fft,
+            window=window,
+            center=False,  # ONNX does not pad
+            onesided=onesided_bool,
+            return_complex=True,
+        )
+
+        # result shape: [batch, bins, frames] (complex)
+        # ONNX expects: [batch, frames, bins, 2]
+
+        # Permute from [batch, bins, frames] to [batch, frames, bins]
+        result = result.permute(0, 2, 1)
+
+        # Convert complex to real representation [batch, frames, bins, 2]
+        result = torch.view_as_real(result)
+
+        return result
+
+    return builder.call_function(
+        _stft,
+        args=(signal, frame_step, window, frame_length, onesided),
+    )
+
+
+# =============================================================================
+# MelWeightMatrix operator
+# =============================================================================
+
+
+@register("MelWeightMatrix", since_version=17)
+def mel_weight_matrix(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Generate a MelWeightMatrix for mel spectrogram computation.
+
+    This operator generates a weight matrix that can be used to convert a linearly
+    sampled frequency spectra (from DFT or STFT) into mel-scaled frequency bins.
+
+    The mel scale is defined as: mel(f) = 2595 * log10(1 + f/700)
+
+    Inputs:
+        num_mel_bins: The number of bands in the mel spectrum (scalar, int32/int64)
+        dft_length: The size of the original DFT (scalar, int32/int64)
+        sample_rate: Samples per second of the input signal (scalar, int32/int64)
+        lower_edge_hertz: Lower bound frequency for mel spectrum (scalar, float)
+        upper_edge_hertz: Upper bound frequency for mel spectrum (scalar, float)
+
+    Attributes:
+        output_datatype: The data type of the output tensor (default: 1 = FLOAT)
+
+    Output:
+        The Mel Weight Matrix with shape [floor(dft_length/2) + 1, num_mel_bins]
+    """
+    from ..utils.dtype import onnx_dtype_to_torch
+
+    num_mel_bins = builder.get_value(node.input[0])
+    dft_length = builder.get_value(node.input[1])
+    sample_rate = builder.get_value(node.input[2])
+    lower_edge_hertz = builder.get_value(node.input[3])
+    upper_edge_hertz = builder.get_value(node.input[4])
+
+    # Get output data type (default is 1 = FLOAT)
+    output_datatype = get_attribute(node, "output_datatype", 1)
+    output_dtype = onnx_dtype_to_torch(output_datatype)
+    if output_dtype is None:
+        output_dtype = torch.float32
+
+    def _mel_weight_matrix(
+        num_mel_bins, dft_length, sample_rate, lower_edge_hertz, upper_edge_hertz, dtype
+    ):
+        # Convert inputs to Python scalars
+        n_mels = int(
+            num_mel_bins.item()
+            if isinstance(num_mel_bins, torch.Tensor)
+            else num_mel_bins
+        )
+        n_fft = int(
+            dft_length.item() if isinstance(dft_length, torch.Tensor) else dft_length
+        )
+        sr = int(
+            sample_rate.item() if isinstance(sample_rate, torch.Tensor) else sample_rate
+        )
+        f_min = float(
+            lower_edge_hertz.item()
+            if isinstance(lower_edge_hertz, torch.Tensor)
+            else lower_edge_hertz
+        )
+        f_max = float(
+            upper_edge_hertz.item()
+            if isinstance(upper_edge_hertz, torch.Tensor)
+            else upper_edge_hertz
+        )
+
+        # Number of spectrogram bins (one-sided DFT)
+        num_spectrogram_bins = n_fft // 2 + 1
+
+        # Create frequency bin indices (n_mels + 2 points for n_mels triangular filters)
+        frequency_bins = torch.arange(0, n_mels + 2, dtype=torch.float32)
+
+        # Convert edge frequencies to mel scale
+        low_frequency_mel = 2595.0 * torch.log10(torch.tensor(1.0 + f_min / 700.0))
+        high_frequency_mel = 2595.0 * torch.log10(torch.tensor(1.0 + f_max / 700.0))
+
+        # Calculate mel step
+        mel_step = (high_frequency_mel - low_frequency_mel) / (frequency_bins.shape[0])
+
+        # Convert to mel frequencies
+        frequency_bins = frequency_bins * mel_step + low_frequency_mel
+
+        # Convert mel frequencies back to Hz
+        frequency_bins = 700.0 * (
+            torch.pow(torch.tensor(10.0), frequency_bins / 2595.0) - 1.0
+        )
+
+        # Convert Hz frequencies to FFT bin indices
+        frequency_bins = ((n_fft + 1) * frequency_bins) // sr
+        frequency_bins = frequency_bins.to(torch.int64)
+
+        # Create the filterbank matrix
+        output = torch.zeros(num_spectrogram_bins, n_mels, dtype=torch.float32)
+
+        for i in range(n_mels):
+            lower_frequency_value = frequency_bins[i].item()  # left
+            center_frequency_point = frequency_bins[i + 1].item()  # center
+            higher_frequency_point = frequency_bins[i + 2].item()  # right
+
+            low_to_center = center_frequency_point - lower_frequency_value
+            if low_to_center == 0:
+                output[center_frequency_point, i] = 1.0
+            else:
+                for j in range(lower_frequency_value, center_frequency_point + 1):
+                    output[j, i] = float(j - lower_frequency_value) / float(
+                        low_to_center
+                    )
+
+            center_to_high = higher_frequency_point - center_frequency_point
+            if center_to_high > 0:
+                for j in range(center_frequency_point, higher_frequency_point):
+                    output[j, i] = float(higher_frequency_point - j) / float(
+                        center_to_high
+                    )
+
+        return output.to(dtype)
+
+    return builder.call_function(
+        _mel_weight_matrix,
+        args=(
+            num_mel_bins,
+            dft_length,
+            sample_rate,
+            lower_edge_hertz,
+            upper_edge_hertz,
+            output_dtype,
+        ),
+    )
+
+
+# =============================================================================
+# Window function operators
+# =============================================================================
+
+
+@register("HannWindow", since_version=17)
+def hann_window(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Generate a Hann window.
+
+    Attributes:
+        periodic: If 1, returns periodic window. If 0, returns symmetric window.
+        output_datatype: ONNX TensorProto data type for output (default: FLOAT).
+    """
+    from ..utils.dtype import onnx_dtype_to_torch
+
+    size = builder.get_value(node.input[0])
+    periodic = get_attribute(node, "periodic", 1)
+    output_datatype = get_attribute(node, "output_datatype", 1)  # Default: FLOAT
+
+    dtype = onnx_dtype_to_torch(output_datatype)
+
+    def _hann_window(
+        window_length: torch.Tensor, periodic: bool, dtype: torch.dtype
+    ) -> torch.Tensor:
+        length = int(window_length.item())
+        return torch.hann_window(length, periodic=periodic, dtype=dtype)
+
+    return builder.call_function(_hann_window, args=(size, bool(periodic), dtype))
+
+
+@register("HammingWindow", since_version=17)
+def hamming_window(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Generate a Hamming window.
+
+    Attributes:
+        periodic: If 1, returns periodic window. If 0, returns symmetric window.
+        output_datatype: ONNX TensorProto data type for output (default: FLOAT).
+
+    Note:
+        ONNX uses specific Hamming coefficients (alpha=0.543478, beta=0.456522)
+        which differ from PyTorch's defaults (0.54, 0.46).
+    """
+    from ..utils.dtype import onnx_dtype_to_torch
+
+    size = builder.get_value(node.input[0])
+    periodic = get_attribute(node, "periodic", 1)
+    output_datatype = get_attribute(node, "output_datatype", 1)  # Default: FLOAT
+
+    dtype = onnx_dtype_to_torch(output_datatype)
+
+    # ONNX HammingWindow uses these specific coefficients
+    alpha = 0.543478
+    beta = 0.456522
+
+    def _hamming_window(
+        window_length: torch.Tensor, periodic: bool, dtype: torch.dtype
+    ) -> torch.Tensor:
+        length = int(window_length.item())
+        return torch.hamming_window(
+            length, periodic=periodic, alpha=alpha, beta=beta, dtype=dtype
+        )
+
+    return builder.call_function(_hamming_window, args=(size, bool(periodic), dtype))
+
+
+@register("BlackmanWindow", since_version=17)
+def blackman_window(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Generate a Blackman window.
+
+    Attributes:
+        periodic: If 1, returns periodic window. If 0, returns symmetric window.
+        output_datatype: ONNX TensorProto data type for output (default: FLOAT).
+    """
+    from ..utils.dtype import onnx_dtype_to_torch
+
+    size = builder.get_value(node.input[0])
+    periodic = get_attribute(node, "periodic", 1)
+    output_datatype = get_attribute(node, "output_datatype", 1)  # Default: FLOAT
+
+    dtype = onnx_dtype_to_torch(output_datatype)
+
+    def _blackman_window(
+        window_length: torch.Tensor, periodic: bool, dtype: torch.dtype
+    ) -> torch.Tensor:
+        length = int(window_length.item())
+        return torch.blackman_window(length, periodic=periodic, dtype=dtype)
+
+    return builder.call_function(_blackman_window, args=(size, bool(periodic), dtype))
+
+
+# =============================================================================
+# Non-maximum suppression
+# =============================================================================
+
+
+@register("NonMaxSuppression")
+def nms(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """Non-maximum suppression for object detection."""
+    boxes = builder.get_value(node.input[0])
+    scores = builder.get_value(node.input[1])
+
+    max_output = get_optional_input(builder, node, 2)
+    iou_threshold = get_optional_input(builder, node, 3)
+    score_threshold = get_optional_input(builder, node, 4)
+
+    # Set defaults if not provided
+    if iou_threshold is None:
+        iou_threshold = 0.0
+    if score_threshold is None:
+        score_threshold = float("-inf")
+
+    center_point_box = get_attribute(node, "center_point_box", 0)
+
+    def _nms(
+        boxes, scores, max_output, iou_threshold, score_threshold, center_point_box
+    ):
+        from torchvision.ops import nms as tv_nms
+
+        batch_size = boxes.shape[0]
+        num_classes = scores.shape[1]
+
+        iou_th = (
+            iou_threshold.item()
+            if isinstance(iou_threshold, torch.Tensor)
+            else iou_threshold
+        )
+        score_th = (
+            score_threshold.item()
+            if isinstance(score_threshold, torch.Tensor)
+            else score_threshold
+        )
+        max_out = (
+            max_output.item()
+            if isinstance(max_output, torch.Tensor) and max_output is not None
+            else max_output
+        )
+
+        results = []
+        for batch_idx in range(batch_size):
+            batch_boxes = boxes[batch_idx]  # [num_boxes, 4]
+
+            # Convert center format to corner format if needed
+            if center_point_box:
+                cx, cy, w, h = (
+                    batch_boxes[:, 0],
+                    batch_boxes[:, 1],
+                    batch_boxes[:, 2],
+                    batch_boxes[:, 3],
+                )
+                batch_boxes = torch.stack(
+                    [cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], dim=1
+                )
+
+            for class_idx in range(num_classes):
+                class_scores = scores[batch_idx, class_idx]  # [num_boxes]
+
+                # Filter by score threshold
+                mask = class_scores > score_th
+                if not mask.any():
+                    continue
+
+                filtered_boxes = batch_boxes[mask]
+                filtered_scores = class_scores[mask]
+
+                # Apply NMS
+                keep = tv_nms(filtered_boxes, filtered_scores, iou_th)
+
+                if max_out is not None:
+                    keep = keep[: int(max_out)]
+
+                # Get original indices
+                original_indices = torch.where(mask)[0][keep]
+
+                for idx in original_indices:
+                    results.append([batch_idx, class_idx, idx.item()])
+
+        if len(results) == 0:
+            return torch.zeros((0, 3), dtype=torch.int64)
+        return torch.tensor(results, dtype=torch.int64)
+
+    return builder.call_function(
+        _nms,
+        args=(
+            boxes,
+            scores,
+            max_output,
+            iou_threshold,
+            score_threshold,
+            center_point_box,
+        ),
+    )

onnx2fx/ops/string.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""String operators.
+
+This module implements ONNX operators for string processing.
+"""
+
+from typing import TYPE_CHECKING
+
+import onnx
+import torch
+
+from ..op_registry import register
+
+if TYPE_CHECKING:
+    from ..graph_builder import GraphBuilder
+
+
+# =============================================================================
+# String operators
+# =============================================================================
+
+
+@register("StringNormalizer")
+def string_normalizer(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """String normalization for numpy string arrays.
+
+    Applies case normalization and stopword filtering to string arrays.
+    """
+    from ..utils.attributes import get_attribute
+
+    x = builder.get_value(node.input[0])
+    case_change_action = get_attribute(node, "case_change_action", "NONE")
+    is_case_sensitive = get_attribute(node, "is_case_sensitive", 0)
+    stopwords = get_attribute(node, "stopwords", [])
+    locale = get_attribute(node, "locale", "")
+
+    def _string_normalizer(
+        arr, case_action: str, case_sensitive: int, stops: list, loc: str
+    ):
+        import numpy as np
+
+        original_shape = arr.shape
+        is_2d = len(original_shape) == 2
+
+        if is_2d:
+            # Process each row separately for 2D arrays
+            result_rows = []
+            for row_idx in range(original_shape[0]):
+                row = arr[row_idx]
+                # Filter stopwords from this row
+                if stops:
+                    if case_sensitive:
+                        stop_set = set(stops)
+                        filtered = [s for s in row if s not in stop_set]
+                    else:
+                        stop_set_lower = {s.lower() for s in stops}
+                        filtered = [
+                            s
+                            for s in row
+                            if not isinstance(s, str) or s.lower() not in stop_set_lower
+                        ]
+                else:
+                    filtered = list(row)
+
+                # Apply case change
+                if case_action == "LOWER":
+                    filtered = [
+                        s.lower() if isinstance(s, str) else s for s in filtered
+                    ]
+                elif case_action == "UPPER":
+                    filtered = [
+                        s.upper() if isinstance(s, str) else s for s in filtered
+                    ]
+
+                result_rows.append(filtered)
+
+            # Find max length and build output
+            if result_rows:
+                # Pad shorter rows if needed (shouldn't happen for valid ONNX)
+                output = np.array(result_rows, dtype=object)
+                return output
+            return np.array([[]], dtype=object)
+
+        # 1D case
+        flat = arr.flatten()
+
+        # Filter stopwords first (before case change for matching)
+        if stops:
+            if case_sensitive:
+                stop_set = set(stops)
+                flat = np.array([s for s in flat if s not in stop_set], dtype=object)
+            else:
+                stop_set_lower = {s.lower() for s in stops}
+                result = []
+                for s in flat:
+                    if isinstance(s, str):
+                        if s.lower() not in stop_set_lower:
+                            result.append(s)
+                    else:
+                        result.append(s)
+                flat = (
+                    np.array(result, dtype=object)
+                    if result
+                    else np.array([], dtype=object)
+                )
+
+        # Apply case change after filtering
+        if case_action == "LOWER":
+            flat = np.array(
+                [s.lower() if isinstance(s, str) else s for s in flat], dtype=object
+            )
+        elif case_action == "UPPER":
+            flat = np.array(
+                [s.upper() if isinstance(s, str) else s for s in flat], dtype=object
+            )
+
+        # Handle empty result - ONNX spec says return [''] for empty
+        if len(flat) == 0:
+            return np.array([""], dtype=object)
+
+        return flat
+
+    return builder.call_function(
+        _string_normalizer,
+        args=(x, case_change_action, is_case_sensitive, stopwords, locale),
+    )