[파이썬] ComfyUI에 이미지 스펙트럼 노드를 만들어 보았습니다.

Image Spectrum Node (세계 최초 ^^)

 

- 어떤 활용이 가능할진 아직 모르겠지만 업계 최초로 ComfyUI에 입력된 이미지의 푸리에 스펙트럼을 만드는 Node 를 만들어 보았습니다.

- 파이썬 패키지로 torch  와 scipy, pillow, 그리고 numpy 등을 사용했습니다.

- ComfyUI 에 이미지 뒤집기나 이진화, 에지검출 등의 간단한 노드만들기는 쉬우나 주파수도메인에서의 변환이나 필터링 같은걸 만드는 건 꽤나 어렵더군요.. 왜냐하면 이미지데이타를 텐서로 바꾸어서 CUDA 메모리상에서 처리해ㅔ 주어야하기 때문입니다.

- 그래도 한번 만들어 봤다는 ~

---

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
}

- 코드가 조악하지만, 그래도 도움이 되시길 바랍니다~

- ComfyUI 의 Custom_Node 디렉토리에 저 파이썬 스크립트를 저장하고 재실행하면 검색으로 저노드가 잡히게 됩니다.

- 노드에서 유일한 파라미터인 Amplitude 는 진폭의 이득이라고 보시면 됩니다. 키울수록 스펙트럼의 진폭이 커지죠(밝아지죠)

참고:
푸리에 영상처리:
https://www.aladin.co.kr/shop/wproduct.aspx?ItemId=309060931

딥러닝을 위한 푸리에 영상처리