From d73fbcaea025e138cb785bf8479ab8e01f181e10 Mon Sep 17 00:00:00 2001
From: Vladislav Tumko <56307628+vectorvp@users.noreply.github.com>
Date: Fri, 15 Nov 2024 20:42:25 +0400
Subject: [PATCH] refactor: Update naming and remove unused package from
 AmplitudeRescaleTransform (#1732)

---
 lightly/transforms/__init__.py                       |  2 +-
 lightly/transforms/amplitude_rescale_transform.py    |  3 +--
 tests/transforms/test_amplitude_rescale_transform.py | 11 +++--------
 3 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/lightly/transforms/__init__.py b/lightly/transforms/__init__.py
index 5474d7600..037fc78a3 100644
--- a/lightly/transforms/__init__.py
+++ b/lightly/transforms/__init__.py
@@ -9,7 +9,7 @@
 # All Rights Reserved
 
 from lightly.transforms.aim_transform import AIMTransform
-from lightly.transforms.amplitude_rescale_transform import AmplitudeRescaleTranform
+from lightly.transforms.amplitude_rescale_transform import AmplitudeRescaleTransform
 from lightly.transforms.byol_transform import (
     BYOLTransform,
     BYOLView1Transform,
diff --git a/lightly/transforms/amplitude_rescale_transform.py b/lightly/transforms/amplitude_rescale_transform.py
index e09128ced..440976aa7 100644
--- a/lightly/transforms/amplitude_rescale_transform.py
+++ b/lightly/transforms/amplitude_rescale_transform.py
@@ -1,12 +1,11 @@
 from typing import Tuple
 
-import numpy as np
 import torch
 from torch import Tensor
 from torch.distributions import Uniform
 
 
-class AmplitudeRescaleTranform:
+class AmplitudeRescaleTransform:
     """Implementation of amplitude rescaling transformation.
 
     This transform will rescale the amplitude of the Fourier Spectrum (`freq_image`) of the image and return it.
diff --git a/tests/transforms/test_amplitude_rescale_transform.py b/tests/transforms/test_amplitude_rescale_transform.py
index 0b8d38eaa..8dea41f43 100644
--- a/tests/transforms/test_amplitude_rescale_transform.py
+++ b/tests/transforms/test_amplitude_rescale_transform.py
@@ -1,11 +1,6 @@
-import numpy as np
 import torch
 
-from lightly.transforms import (
-    AmplitudeRescaleTranform,
-    IRFFT2DTransform,
-    RFFT2DTransform,
-)
+from lightly.transforms import AmplitudeRescaleTransform, RFFT2DTransform
 
 
 # Testing function image -> FFT -> AmplitudeRescale.
@@ -16,10 +11,10 @@ def test() -> None:
     rfftTransform = RFFT2DTransform()
     rfft = rfftTransform(image)
 
-    ampRescaleTf_1 = AmplitudeRescaleTranform()
+    ampRescaleTf_1 = AmplitudeRescaleTransform()
     rescaled_rfft_1 = ampRescaleTf_1(rfft)
 
-    ampRescaleTf_2 = AmplitudeRescaleTranform(range=(1.0, 2.0))
+    ampRescaleTf_2 = AmplitudeRescaleTransform(range=(1.0, 2.0))
     rescaled_rfft_2 = ampRescaleTf_2(rfft)
 
     assert rescaled_rfft_1.shape == rfft.shape