[Pytorch] torch.fft 작동 방식

파이토치를 이용해 이미지를 frequency domain으로 변환시켜주는 함수들에 대해 정리해보려 한다. 파이토치의 torch.fft 모듈은 푸리에 변환(Fourier Transform)을 수행하는 다양한 함수를 제공한다. ( 관련 모든 함수들은 다음 링크를 참조 https://pytorch.org/docs/stable/fft.html )

이번 글에서는 torch.fft.fftn, torch.fft.fftshift, torch.fft.rfftn, 그리고 역변환인 torch.fft.ifftn의 작동 방식을 설명하고 실행예시를 정리하려 한다.

torch.fft.fftn

가장 먼저 torch.fft.fftn 이다. n 차원의 tensor에 discrete Fourier transform을 적용시켜주는 함수이다.

parameter로 dim을 넣어줄 수 있는데 넣지 않으면 default로 모든 dimension에 대해 fft를 적용한다. 이미지에 대해 적용한다고 했을때는 H,W가 존재하는 dimension을 paramter로 명시해주면 된다. 여기서는 편의상 이미지를 256*256 형태의 gray scale이라고 하자 그러면,

fft_output = torch.fft.fftn(image)

형태로 간단히 변환할 수 있다.

기본적으로 변환된 fft_output은 이미지와 같은 size로 256*256 형태의 complex 값이 된다. 간단히 visualize하기 위해 magnitude만 따면 아래와 같은 결과를 확인할 수 있다.

torch.fft.fftshift

근데 보통 우리가 이미지를 frequency에서 분석할때는 low frequency 값들이 중앙에 있도록 하는데 기본 torch.fft.fftn은 low frequency가 바깥쪽에 배치된다. 그래서 fft_shifted = torch.fft.fftshift(fft_output) 을 사용하면 low frequency 지역을 가운데로 옮겨줄 수 있다.

torch.fft.rfftn

그리고 frequency domain의 특성상 원점 대칭의 결과를 가진다. 따라서 output을 전체 이미지로 따지 않고 절반만 따도 똑같은 정보량을 가지고 있다. 그래서

rfft_output = torch.fft.rfftn(image)

을 통해 불필요한 결과는 없애 결과를 얻을 수도 있다. 아래는 rfft변환 후 shift까지 적용해서 magnitude를 측정한 결과이다

torch.fft.ifftshift/ torch.fft.ifftn

마지막으로 frequency에서 이미지로 다시 변환하고 싶으면 ifftn을 적용하면 된다. 이 때 만약 frequency에서 shift를 적용한 상태라면 shift의 역변환인 ifftshift를 먼저 적용하고 다시 역변환을 시켜줘야 한다.

reconstructed_img = torch.fft.ifftn(torch.fft.ifftshift(fft_output)).real

위와 같이 적용이 되고 torch.fft.ifftn의 결과는 complex형태로 출력되기 때문에 real 혹은 torch.abs()로 결과를 다시 실수로 만들어줘야 한다. 최종적으로 아래와 같이 원본과 동일한 이미지를 다시 얻어낼 수 있다.