vendored model

sharp/models/normalizers.py

@@ -0,0 +1,80 @@
+"""Contains an implementation of image normalizers for perceptual loss.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence, Union
+
+import torch
+from torch import nn
+
+
+class MeanStdNormalizer(nn.Module):
+    """Normalizing image input by mean and std."""
+
+    mean: torch.Tensor
+    std_inv: torch.Tensor
+
+    def __init__(
+        self,
+        mean: Union[Sequence[float], torch.Tensor],
+        std: Union[Sequence[float], torch.Tensor],
+    ):
+        """Initialize MeanStdNormalizer."""
+        super(MeanStdNormalizer, self).__init__()
+        if not isinstance(mean, torch.Tensor):
+            mean = torch.as_tensor(mean).view(-1, 1, 1)
+        if not isinstance(std, torch.Tensor):
+            std = torch.as_tensor(std).view(-1, 1, 1)
+        self.register_buffer("mean", mean)
+        # We use inverse std to use a multiplication which is better supported by the hardware
+        self.register_buffer("std_inv", 1.0 / std)
+
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        """Apply mean and std normalization over input image."""
+        return (image - self.mean) * self.std_inv
+
+
+class AffineRangeNormalizer(nn.Module):
+    """Perform linear mapping to map input_range to output_range.
+
+    Output_range defaults to (0, 1).
+    """
+
+    def __init__(
+        self,
+        input_range: tuple[float, float],
+        output_range: tuple[float, float] = (0, 1),
+    ):
+        """Initialize AffineRangeNormalizer."""
+        super().__init__()
+        input_min, input_max = input_range
+        output_min, output_max = output_range
+        if input_max <= input_min:
+            raise ValueError(f"Invalid input_range: {input_range}")
+        if output_max <= output_min:
+            raise ValueError(f"Invalid output_range: {output_range}")
+
+        self.scale = (output_max - output_min) / (input_max - input_min)
+        self.bias = output_min - input_min * self.scale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply affine range normalization over input image."""
+        if self.scale != 1.0:
+            x = x * self.scale
+
+        if self.bias != 0.0:
+            x = x + self.bias
+
+        return x
+
+
+class MobileNetNormalizer(AffineRangeNormalizer):
+    """Image normalization in mobilenet."""
+
+    def __init__(self, input_range: tuple[float, float] = (0, 1)):
+        """Initialize MobileNetNormalizer."""
+        super().__init__(input_range=input_range, output_range=(-1, 1))