From 6d25f6e93d634aa3f7ca104857a2513e693a4640 Mon Sep 17 00:00:00 2001
From: Gromit Park <bemore.one@gmail.com>
Date: Fri, 28 Mar 2025 21:09:59 +0900
Subject: [PATCH] Add main node implementation

---
 __init__.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 __init__.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..7f7f0ed
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,43 @@
+import torch
+import numpy as np
+from PIL import Image, ImageOps, ImageFilter
+from scipy.fft import fft2, fftshift
+
+class Image_Spectrum:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "amplitude": ("INT", {
+                    "default": 20,
+                    "min": 1,
+                    "max": 50,
+                    "step": 1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "image_spectrum"
+    CATEGORY = "image"
+
+    def image_spectrum(self, image, amplitude):
+        image = 255. * image[0].cpu().numpy()
+        image = Image.fromarray(np.clip(image, 0, 255).astype(np.uint8))
+        image = ImageOps.grayscale(image)
+        
+        # 2D Fourier 변환 수행
+        f_transform = fft2(image)
+        f_shift = fftshift(f_transform)
+        spectrum = amplitude*np.log1p(np.abs(f_shift))
+        
+        image = np.array(spectrum).astype(np.float32) / 255.0
+        mask = 1. - torch.from_numpy(image)
+        image = torch.from_numpy(image)[None,]
+        
+        return (image,)
+
+NODE_CLASS_MAPPINGS = {
+    "Image_Spectrum": Image_Spectrum
+}