430 lines
16 KiB
Python
430 lines
16 KiB
Python
import torch
|
||
#import torch.nn.functional as F
|
||
import torchvision.transforms.v2 as T
|
||
import numpy as np
|
||
import comfy.utils
|
||
import comfy.model_management as mm
|
||
import gc
|
||
from tqdm import tqdm
|
||
import cv2
|
||
import numpy as np
|
||
from scipy.interpolate import RBFInterpolator
|
||
import folder_paths
|
||
import os
|
||
import logging
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
log = logging.getLogger(__name__)
|
||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
|
||
def tensor_to_image(image):
|
||
return np.array(T.ToPILImage()(image.permute(2, 0, 1)).convert('RGB'))
|
||
|
||
def image_to_tensor(image):
|
||
return T.ToTensor()(image).permute(1, 2, 0)
|
||
#return T.ToTensor()(Image.fromarray(image)).permute(1, 2, 0)
|
||
|
||
|
||
class FFTNode:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", ),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE","FFTData")
|
||
RETURN_NAMES = ("image","FFTData")
|
||
FUNCTION = "toFFT"
|
||
CATEGORY = "FFT"
|
||
def toFFT(self, image):
|
||
FFTImageList=[]
|
||
FFT_Channel_Data=[]
|
||
channelCount=3
|
||
imageCount=image.shape[0]
|
||
# print("image.shape:",image.shape)
|
||
# ## torch.Size([1, 64, 64, 3]) 表示来了1张图，64*64，3通道 loadImage节点过来就是这个
|
||
# print("type of image:",type(image)) ## type of image: <class 'torch.Tensor'>
|
||
# pixel = image[0][0][0]
|
||
# print("pixel:",pixel) ## pixel: tensor([0., 0., 0.])
|
||
# tValue=pixel[0]
|
||
# print("Type of tValue:",type(tValue)) ## type of tensors_out: <class 'torch.Tensor'>
|
||
|
||
for i in range(image.shape[0]):
|
||
sourceImg = image[i]
|
||
cv2Image = (sourceImg.contiguous() * 255).byte()
|
||
#image1=image[0]
|
||
#print(image1[0][0])
|
||
## tensor([0., 0., 0.])
|
||
## 是0-1之间的浮点数
|
||
|
||
dim = sourceImg.dim()
|
||
if dim == 3:
|
||
|
||
#print(cv2Image.shape)
|
||
##torch.Size([64, 64, 3]) 变成0-255的byte, 颜色值, 3通道
|
||
R_channel = cv2Image[:, :,0]
|
||
G_channel = cv2Image[:, :,1]
|
||
B_channel = cv2Image[:, :,2]
|
||
|
||
fshiftData=[]
|
||
#傅里叶变换
|
||
R_fft = np.fft.fft2(R_channel)
|
||
R_fshift = np.fft.fftshift(R_fft)
|
||
fshiftData.append(R_fshift)
|
||
G_fft = np.fft.fft2(G_channel)
|
||
G_fshift = np.fft.fftshift(G_fft)
|
||
fshiftData.append(G_fshift)
|
||
B_fft = np.fft.fft2(B_channel)
|
||
B_fshift = np.fft.fftshift(B_fft)
|
||
fshiftData.append(B_fshift)
|
||
|
||
R_img = np.log(np.abs(R_fshift))
|
||
B_img = np.log(np.abs(B_fshift))
|
||
G_img = np.log(np.abs(G_fshift))
|
||
|
||
R_img = R_img / np.max(R_img)
|
||
B_img = B_img / np.max(B_img)
|
||
G_img = G_img / np.max(G_img)
|
||
|
||
fftImg = np.dstack((R_img, G_img, B_img))
|
||
fftImg = fftImg.astype(np.float32)
|
||
#image = image.astype(np.uint8)
|
||
FFT_Channel_Data.append(fshiftData)
|
||
# # # one_channel = cv2Image[:, :,2]
|
||
#print(red_channel.shape)
|
||
## torch.Size([64, 64]) 拿到其中一个颜色通道
|
||
|
||
# # # f = np.fft.fft2(one_channel)
|
||
# # # #dft = cv2.dft(np.float32(one_channel), flags = cv2.DFT_COMPLEX_OUTPUT)
|
||
# # # FFT_Data.append(f)
|
||
# # # #print(f.shape)
|
||
# # # ## (64, 64) 现在不是torch, tensor了，是个二维数组
|
||
|
||
# # # #默认结果中心点位置是在左上角,
|
||
# # # #调用fftshift()函数转移到中间位置
|
||
# # # fshift = np.fft.fftshift(f)
|
||
|
||
# # # #fft结果是复数, 其绝对值结果是振幅
|
||
# # # fimg = np.log(np.abs(fshift))
|
||
# # # #print(fimg.shape)
|
||
# # # ## (64, 64)
|
||
# # # ## 取复数的模，然后取对数。仅用于显示，原始数据还在f里
|
||
# # # ## 此时是取对数后的float结果，范围不一定是0-255，也不一定是0-1，可能很大也可能很小
|
||
# # # max_v = np.max(fimg)
|
||
# # # fimg = fimg / max_v
|
||
# # # print(max_v)
|
||
|
||
# # # img_3_channels = np.repeat(fimg[:, :, np.newaxis], 3, axis=2)
|
||
#print(img_3_channels.shape)
|
||
## (64, 64, 3) 复制到其余两个通道。让它变成灰度图
|
||
#print(img_3_channels[0][0])
|
||
FFTImageList.append(fftImg)
|
||
else:
|
||
channelCount=1
|
||
fshiftData=[]
|
||
#傅里叶变换
|
||
R_fft = np.fft.fft2(cv2Image)
|
||
R_fshift = np.fft.fftshift(R_fft)
|
||
fshiftData.append(R_fshift)
|
||
fftImg = np.log(np.abs(R_fshift))
|
||
fftImg = fftImg / np.max(fftImg)
|
||
fftImg = fftImg.astype(np.float32)
|
||
FFTImageList.append(fftImg)
|
||
FFT_Channel_Data.append(fshiftData)
|
||
|
||
tensors_out = (
|
||
torch.stack([torch.from_numpy(np_array) for np_array in FFTImageList])
|
||
)
|
||
FFT_Data={'channelCount':channelCount,'FFT_Channel_Data':FFT_Channel_Data,'imageCount':imageCount}
|
||
|
||
# print("in FFTNode, before output")
|
||
# print("type of tensors_out:",type(tensors_out))
|
||
# print("tensors_out shape:",tensors_out.shape)
|
||
# print("type of tensors_out[0]:",type(tensors_out[0]))
|
||
# print("tensors_out shape[0]:",tensors_out[0].shape)
|
||
# pixel = tensors_out[0][0][0]
|
||
# print("pixel:",pixel) ## pixel: tensor([0., 0., 0.])
|
||
# tValue=pixel[0]
|
||
# print("Type of tValue:",type(tValue)) ## Type of tValue: <class 'torch.Tensor'>
|
||
return (tensors_out,FFT_Data)
|
||
|
||
def ApplyMaskAndInvert(l_fshift,l_mask):
|
||
# print("in fucntion, mask shape:",l_mask.shape)
|
||
# print("in fucntion, fshift shape:",l_fshift.shape)
|
||
f = l_fshift * l_mask
|
||
ishift = np.fft.ifftshift(f)
|
||
iimg = np.fft.ifft2(ishift)
|
||
iimg = np.abs(iimg)
|
||
return f,iimg
|
||
|
||
class InvertFFTNode:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"ff": ("FFTData", ),
|
||
"filterRadius": ("INT", {
|
||
"default": 100,
|
||
})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE","IMAGE","IMAGE",)
|
||
RETURN_NAMES = ("image","IMAGE","IMAGE",)
|
||
FUNCTION = "fromFFT"
|
||
CATEGORY = "FFT"
|
||
def DoOneChannel(self,fshift,filterRadius):
|
||
FFTImageList=[]
|
||
#f=ff[0][0]
|
||
#fshift = np.fft.fftshift(f)
|
||
#傅里叶逆变换
|
||
#ishift = np.fft.ifftshift(fshift)
|
||
iimg = np.fft.ifft2(fshift)
|
||
iimg = np.abs(iimg)
|
||
FFTImageList.append(iimg)
|
||
|
||
#设置高通滤波器
|
||
rows, cols = fshift.shape
|
||
crow,ccol = int(rows/2), int(cols/2)
|
||
mask = np.ones((rows, cols))
|
||
mask[crow-filterRadius:crow+filterRadius, ccol-filterRadius:ccol+filterRadius] = 0
|
||
#fshift[crow-30:crow+30, ccol-30:ccol+30] = 0
|
||
|
||
mask2 = np.zeros((rows, cols))
|
||
mask2[crow-filterRadius:crow+filterRadius, ccol-filterRadius:ccol+filterRadius] = 1
|
||
|
||
#傅里叶逆变换
|
||
fshiftmask,hiPassImg = ApplyMaskAndInvert(fshift,mask)
|
||
FFTImageList.append(mask*255)
|
||
# # ishift = np.fft.ifftshift(fshift)
|
||
# # iimg = np.fft.ifft2(ishift)
|
||
# # iimg = np.abs(iimg)
|
||
# print("np.max(hiPass):",np.max(hiPassImg))
|
||
# print("np.min(hiPass):",np.min(hiPassImg))
|
||
FFTImageList.append(hiPassImg)
|
||
|
||
fimg = np.log(np.abs(fshiftmask))
|
||
#print(fimg.shape)
|
||
## (64, 64)
|
||
## 取复数的模，然后取对数。仅用于显示，原始数据还在f里
|
||
## 此时是取对数后的float结果，范围不一定是0-255，也不一定是0-1，可能很大也可能很小
|
||
max_v = np.max(fimg)
|
||
fimg = fimg / max_v * 254
|
||
# print("max_v of fimg:",max_v)
|
||
# print("max_v of after fimg:",np.max(fimg))
|
||
|
||
FFTImageList.append(fimg)
|
||
fshiftmask,lowPassImg = ApplyMaskAndInvert(fshift,mask2)
|
||
FFTImageList.append(mask2*255)
|
||
# print("np.max(lowPassImg):",np.max(lowPassImg))
|
||
# print("np.min(lowPassImg):",np.min(lowPassImg))
|
||
fimg = np.log(np.abs(fshiftmask))
|
||
#print(fimg.shape)
|
||
## (64, 64)
|
||
## 取复数的模，然后取对数。仅用于显示，原始数据还在f里
|
||
## 此时是取对数后的float结果，范围不一定是0-255，也不一定是0-1，可能很大也可能很小
|
||
max_v = np.max(fimg)
|
||
fimg = fimg / max_v * 254
|
||
# print("max_v of fimg:",max_v)
|
||
# print("max_v of after fimg:",np.max(fimg))
|
||
#img_3_channels = np.repeat(fimg[:, :, np.newaxis], 3, axis=2)
|
||
FFTImageList.append(fimg)
|
||
FFTImageList.append(lowPassImg)
|
||
|
||
tensors_out = (
|
||
torch.stack([torch.from_numpy(np_array) for np_array in FFTImageList])/255
|
||
)
|
||
return tensors_out
|
||
def fromFFT(self, ff,filterRadius):
|
||
f0=ff[0][0]
|
||
#f1=ff[0][1]
|
||
#f2=ff[0][2]
|
||
out1=self.DoOneChannel(f0,filterRadius)
|
||
#print("C1")
|
||
#out2=self.DoOneChannel(f1,filterRadius)
|
||
#print("C2")
|
||
#out3=self.DoOneChannel(f2,filterRadius)
|
||
return (out1,out1,out1,)
|
||
|
||
|
||
class InvertFFTWithMask:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"ff": ("FFTData", ),
|
||
"mask": ("MASK",)
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
FUNCTION = "fromFFT"
|
||
CATEGORY = "FFT"
|
||
def DoOneChannel(self,fshift,mask):
|
||
fshiftmask,hiPassImg = ApplyMaskAndInvert(fshift,mask)
|
||
return hiPassImg/255
|
||
def fromFFT(self, ff,mask):
|
||
mask0=mask[0]
|
||
maskNP=np.array(mask0)
|
||
channelCount = ff['channelCount']
|
||
imageCount = ff['imageCount']
|
||
### ff['FFT_Channel_Data'] 是数组，每个元素对应一张图片。每个成员也是数组。
|
||
### ff['FFT_Channel_Data'][0] 是一个数组，对应一张图片。成员数量是3或者1，对应3个通道或者1个通道
|
||
res=[]
|
||
for i in range(imageCount):
|
||
if channelCount==3:
|
||
f0 = ff['FFT_Channel_Data'][i][0]
|
||
#print("type of f0:",type(f0))
|
||
out0=self.DoOneChannel(f0,maskNP)
|
||
#print("type of out0:",type(out0))
|
||
#print("out0 shape:",out0.shape)
|
||
f1 = ff['FFT_Channel_Data'][i][1]
|
||
out1=self.DoOneChannel(f1,maskNP)
|
||
#print("type of out1:",type(out1))
|
||
#print("out1 shape:",out1.shape)
|
||
f2 = ff['FFT_Channel_Data'][i][2]
|
||
out2=self.DoOneChannel(f2,maskNP)
|
||
#print("type of out2:",type(out2))
|
||
#print("out2 shape:",out2.shape)
|
||
doneImg = np.dstack((out0, out1, out2))
|
||
doneImg = doneImg.astype(np.float32)
|
||
#print("doneImg shape:",doneImg.shape)
|
||
res.append(doneImg)
|
||
else:
|
||
f0 = ff['FFT_Channel_Data'][i]
|
||
doneImg=self.DoOneChannel(f0,maskNP)
|
||
doneImg = doneImg.astype(np.float32)
|
||
#doneImg = out0
|
||
res.append(doneImg)
|
||
#print("C1")
|
||
#out2=self.DoOneChannel(f1,filterRadius)
|
||
#print("C2")
|
||
#out3=self.DoOneChannel(f2,filterRadius)
|
||
tensors_out = (
|
||
torch.stack([torch.from_numpy(np_array) for np_array in res])
|
||
)
|
||
# print("in InvertFFTWitMask, before output")
|
||
# print("type of tensors_out:",type(tensors_out))
|
||
# print("tensors_out shape:",tensors_out.shape)
|
||
|
||
# print("type of tensors_out[0]:",type(tensors_out[0]))
|
||
# print("tensors_out shape[0]:",tensors_out[0].shape)
|
||
|
||
# pixel = tensors_out[0][0][0]
|
||
# print("pixel shape:",pixel.shape)
|
||
# print("pixel:",pixel)
|
||
# tvalue=pixel[0]
|
||
# print("tvalue:",tvalue)
|
||
# print("tvalue type:",type(tvalue))
|
||
|
||
return (tensors_out,)
|
||
|
||
class FindFFTSpot:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
#"FFTImage": ("IMAGE", ),
|
||
"ff": ("FFTData", ),
|
||
"hiPassFactor": ("INT", {"default": 10, "min": 1, "max": 1024}),
|
||
"centerRange": ("INT", {"default": 30, "min": 1, "max": 1024}),
|
||
"crossWidth": ("INT", {"default": 30, "min": 0, "max": 100}),
|
||
"threshold": ("INT", {"default": 100, "min": 0, "max": 254}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
FUNCTION = "FindSpot"
|
||
CATEGORY = "FFT"
|
||
def FindSpot(self, ff,hiPassFactor,centerRange,crossWidth,threshold):
|
||
FFTImageList=[]
|
||
#imageCount=FFTImage.shape[0]
|
||
|
||
channelCount = ff['channelCount']
|
||
imageCount = ff['imageCount']
|
||
|
||
for i in range(imageCount):
|
||
|
||
if(channelCount==1):
|
||
ffData = ff['FFT_Channel_Data'][i]
|
||
else:
|
||
ffData = ff['FFT_Channel_Data'][i][0]
|
||
|
||
rows, cols = ffData.shape
|
||
crow,ccol = int(rows/2), int(cols/2) #中心位置
|
||
mask = np.ones((rows,cols),np.uint8)
|
||
mask[crow-hiPassFactor:crow+hiPassFactor, ccol-hiPassFactor:ccol+hiPassFactor] = 0
|
||
|
||
_,R_channel = ApplyMaskAndInvert(ffData,mask)
|
||
|
||
|
||
|
||
|
||
# sourceImg=FFTImage[i]
|
||
# cv2Image = (sourceImg.contiguous() * 255).byte()
|
||
# dim = channelCount
|
||
# if dim == 3:
|
||
# R_channel = cv2Image[:, :,0]
|
||
# else:
|
||
# R_channel = cv2Image
|
||
# #print("R_channel shape:",R_channel.shape)
|
||
# R_channel = np.asanyarray(R_channel)
|
||
# #print("R_channel AS NP shape:",R_channel.shape)
|
||
|
||
min_val, max_val = np.min(R_channel), np.max(R_channel)
|
||
out1 = (R_channel - min_val)*(255.0/(max_val - min_val))
|
||
|
||
rows, cols = R_channel.shape
|
||
crow,ccol = int(rows/2), int(cols/2) #中心位置
|
||
|
||
out1[crow-centerRange:crow+centerRange, ccol-centerRange:ccol+centerRange] = 0
|
||
out1[:, ccol-crossWidth:ccol+crossWidth] = 0
|
||
out1[crow-crossWidth:crow+crossWidth, :] = 0
|
||
if(threshold>0):
|
||
out1[out1<threshold]=0
|
||
out1[out1>threshold]=255
|
||
fftImg = np.dstack((out1, out1, out1))
|
||
fftImg = fftImg.astype(np.float32)
|
||
FFTImageList.append(fftImg)
|
||
|
||
tensors_out = (
|
||
torch.stack([torch.from_numpy(np_array) for np_array in FFTImageList])/255
|
||
)
|
||
# print("find spot")
|
||
# print("type of tensors_out:",type(tensors_out))
|
||
# print("tensors_out shape:",tensors_out.shape)
|
||
|
||
# print("type of tensors_out[0]:",type(tensors_out[0]))
|
||
# print("tensors_out shape[0]:",tensors_out[0].shape)
|
||
|
||
# pixel = tensors_out[0][0][0]
|
||
# print("pixel:",pixel) ## pixel: tensor([0., 0., 0.])
|
||
# tValue=pixel[0]
|
||
# print("Type of tValue:",type(tValue)) ## Type of tValue: <class 'float'>
|
||
|
||
return (tensors_out,)
|
||
|
||
|
||
|
||
# A dictionary that contains all nodes you want to export with their names
|
||
# NOTE: names should be globally unique
|
||
NODE_CLASS_MAPPINGS = {
|
||
"InvertFFTNode":InvertFFTNode,
|
||
"FFTNode":FFTNode,
|
||
"InvertFFTWithMask":InvertFFTWithMask,
|
||
"FindFFTSpot":FindFFTSpot,
|
||
}
|
||
|
||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"InvertFFTNode":"InvertFFTNode",
|
||
"FFTNode":"FFTNode",
|
||
"InvertFFTWithMask":"InvertFFTWithMask",
|
||
"FindFFTSpot":"FindFFTSpot",
|
||
|
||
}
|