more tonemap options, fixed constant color, added normal map converter and batch average image

This commit is contained in:
spacepxl
2024-02-12 03:10:05 -05:00
parent 0b7b12382b
commit 46d32933e8
2 changed files with 160 additions and 27 deletions
+12 -8
View File
@@ -66,13 +66,13 @@ Normalizes each frame in a batch to the overall mean and std dev, good for remov
Absolute value of the difference between inputs, with a multiplier to boost dark values for easier viewing. Alternative to the vanilla merge difference node, which is only a subtraction without the abs()
### Image Constant
### Image Constant (RGB/HSV)
Creates an empty image of any color, with a color picker UI.
Create an empty image of any color, either RGB or HSV
### Offset Latent Image
Creates an empty latent image with custom values, for offset noise but with per-channel control.
Create an empty latent image with custom values, for offset noise but with per-channel control. Can be combined with Latent Stats to get channel values.
### Latent Stats
@@ -108,9 +108,13 @@ b = b > 0.06 ? pow(10, (SB * 1023 - 420)/261.5) * 0.19 - 0.01 : b
### Exposure Adjust
Linear exposure adjustment in f-stops
Linear exposure adjustment in f-stops, with optional tonemap.
### Convert Normals
Translate between different normal map color spaces, with optional normalization fix and black region fix.
### Batch Average Image
Returns the single average image of a batch.
## TODO:
- bilateral filter image for single frame denoise
- temporal filters for video denoise
- deconvolution
+148 -19
View File
@@ -647,20 +647,63 @@ class ImageConstant:
return {"required": { "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
"color": ("COLOR",),
"red": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"green": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"blue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "image"
CATEGORY = "image/filters"
def generate(self, width, height, batch_size, color):
# rgb = color.lstrip('#')
rgb = tuple(int(color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
def generate(self, width, height, batch_size, red, green, blue):
r = torch.full([batch_size, height, width, 1], red)
g = torch.full([batch_size, height, width, 1], green)
b = torch.full([batch_size, height, width, 1], blue)
return (torch.cat((r, g, b), dim=-1), )
def hsv_to_rgb(h, s, v):
if s:
if h == 1.0: h = 0.0
i = int(h*6.0)
f = h*6.0 - i
r = torch.full([batch_size, height, width, 1], rgb[0]/255)
g = torch.full([batch_size, height, width, 1], rgb[1]/255)
b = torch.full([batch_size, height, width, 1], rgb[2]/255)
w = v * (1.0 - s)
q = v * (1.0 - s * f)
t = v * (1.0 - s * (1.0 - f))
if i==0: return (v, t, w)
if i==1: return (q, v, w)
if i==2: return (w, v, t)
if i==3: return (w, q, v)
if i==4: return (t, w, v)
if i==5: return (v, w, q)
else: return (v, v, v)
class ImageConstantHSV:
def __init__(self, device="cpu"):
self.device = device
@classmethod
def INPUT_TYPES(s):
return {"required": { "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
"hue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"saturation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "image/filters"
def generate(self, width, height, batch_size, hue, saturation, value):
red, green, blue = hsv_to_rgb(hue, saturation, value)
r = torch.full([batch_size, height, width, 1], red)
g = torch.full([batch_size, height, width, 1], green)
b = torch.full([batch_size, height, width, 1], blue)
return (torch.cat((r, g, b), dim=-1), )
class OffsetLatentImage:
@@ -695,8 +738,8 @@ class LatentStats:
def INPUT_TYPES(s):
return {"required": {"latent": ("LATENT", ),}}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("stats",)
RETURN_TYPES = ("STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
RETURN_NAMES = ("stats", "c0_mean", "c1_mean", "c2_mean", "c3_mean")
FUNCTION = "notify"
OUTPUT_NODE = True
@@ -711,10 +754,12 @@ class LatentStats:
text.append(f"width: {width} ({width * 8})")
text.append(f"height: {height} ({height * 8})")
cmean = [0,0,0,0]
for i in range(4):
minimum = torch.min(latents[:,i,:,:]).item()
maximum = torch.max(latents[:,i,:,:]).item()
std_dev, mean = torch.std_mean(latents[:,i,:,:], dim=None)
cmean[i] = mean
text.append(f"c{i} mean: {mean:.1f} std_dev: {std_dev:.1f} min: {minimum:.1f} max: {maximum:.1f}")
@@ -730,7 +775,7 @@ class LatentStats:
returntext += text[i]
print(printtext)
return (returntext,)
return (returntext, cmean[0], cmean[1], cmean[2], cmean[3])
def sRGBtoLinear(npArray):
less = npArray <= 0.0404482362771082
@@ -742,16 +787,20 @@ def linearToSRGB(npArray):
npArray[less] = npArray[less] * 12.92
npArray[~less] = np.power(npArray[~less], 1/2.4) * 1.055 - 0.055
def linearToTonemap(npArray):
def linearToTonemap(npArray, tonemap_scale):
npArray /= tonemap_scale
more = npArray > 0.06
SLog3 = np.clip((np.log10((npArray + 0.01)/0.19) * 261.5 + 420) / 1023, 0, 1)
npArray[more] = np.power(1 / (1 + (1 / np.power(SLog3[more] / (1 - SLog3[more]), 1.7))), 1.7)
npArray *= tonemap_scale
def tonemapToLinear(npArray):
def tonemapToLinear(npArray, tonemap_scale):
npArray /= tonemap_scale
more = npArray > 0.06
x = np.power(np.clip(npArray, 0, 1), 1/1.7)
x = np.power(np.clip(npArray, 0.000001, 1), 1/1.7)
ut = 1 / (1 + np.power((-1 / x) * (x - 1), 1/1.7))
npArray[more] = np.power(10, (ut[more] * 1023 - 420)/261.5) * 0.19 - 0.01
npArray *= tonemap_scale
class Tonemap:
@classmethod
@@ -761,6 +810,7 @@ class Tonemap:
"images": ("IMAGE",),
"input_mode": (["linear", "sRGB"],),
"output_mode": (["sRGB", "linear"],),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
@@ -769,13 +819,13 @@ class Tonemap:
CATEGORY = "image/filters"
def apply(self, images, input_mode, output_mode):
def apply(self, images, input_mode, output_mode, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[:,:,:,:3])
linearToTonemap(t[:,:,:,:3])
linearToTonemap(t[:,:,:,:3], tonemap_scale)
if output_mode == "sRGB":
linearToSRGB(t[:,:,:,:3])
@@ -792,6 +842,7 @@ class UnTonemap:
"images": ("IMAGE",),
"input_mode": (["sRGB", "linear"],),
"output_mode": (["linear", "sRGB"],),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
@@ -800,13 +851,13 @@ class UnTonemap:
CATEGORY = "image/filters"
def apply(self, images, input_mode, output_mode):
def apply(self, images, input_mode, output_mode, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[:,:,:,:3])
tonemapToLinear(t[:,:,:,:3])
tonemapToLinear(t[:,:,:,:3], tonemap_scale)
if output_mode == "sRGB":
linearToSRGB(t[:,:,:,:3])
@@ -828,6 +879,8 @@ class ExposureAdjust:
"stops": ("FLOAT", {"default": 0.0, "min": -100, "max": 100, "step": 0.01}),
"input_mode": (["sRGB", "linear"],),
"output_mode": (["sRGB", "linear"],),
"use_tonemap": ("BOOLEAN", {"default": False}),
"tonemap_scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 10, "step": 0.01}),
},
}
@@ -836,14 +889,20 @@ class ExposureAdjust:
CATEGORY = "image/filters"
def apply(self, images, stops, input_mode, output_mode):
def apply(self, images, stops, input_mode, output_mode, use_tonemap, tonemap_scale):
t = images.detach().clone().cpu().numpy().astype(np.float32)
if input_mode == "sRGB":
sRGBtoLinear(t[:,:,:,:3])
if use_tonemap:
tonemapToLinear(t[:,:,:,:3], tonemap_scale)
exposure(t[:,:,:,:3], stops)
if use_tonemap:
linearToTonemap(t[:,:,:,:3], tonemap_scale)
if output_mode == "sRGB":
linearToSRGB(t[:,:,:,:3])
t = np.clip(t, 0, 1)
@@ -851,6 +910,70 @@ class ExposureAdjust:
t = torch.from_numpy(t)
return (t,)
# Normal map standard coordinates: +r:+x:right, +g:+y:up, +b:+z:in
class ConvertNormals:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"normals": ("IMAGE",),
"input_mode": (["BAE", "MiDaS", "Standard"],),
"output_mode": (["BAE", "MiDaS", "Standard"],),
"scale_XY": ("FLOAT",{"default": 1, "min": 0, "max": 100, "step": 0.001}),
"normalize": ("BOOLEAN", {"default": True}),
"fix_black": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "convert_normals"
CATEGORY = "image/filters"
def convert_normals(self, normals, input_mode, output_mode, scale_XY, normalize, fix_black):
t = normals.detach().clone()
if input_mode == "BAE":
t[:,:,:,0] = 1 - t[:,:,:,0] # invert R
elif input_mode == "MiDaS":
t[:,:,:,:3] = torch.stack([1 - t[:,:,:,2], t[:,:,:,1], t[:,:,:,0]], dim=3) # BGR -> RGB and invert R
if fix_black:
key = torch.clamp(1 - t[:,:,:,2] * 2, min=0, max=1)
t[:,:,:,0] += key * 0.5
t[:,:,:,1] += key * 0.5
t[:,:,:,2] += key
t[:,:,:,:2] = (t[:,:,:,:2] - 0.5) * scale_XY + 0.5
if normalize:
t[:,:,:,:3] = torch.nn.functional.normalize(t[:,:,:,:3] * 2 - 1, dim=3) / 2 + 0.5
if output_mode == "BAE":
t[:,:,:,0] = 1 - t[:,:,:,0] # invert R
elif output_mode == "MiDaS":
t[:,:,:,:3] = torch.stack([t[:,:,:,2], t[:,:,:,1], 1 - t[:,:,:,0]], dim=3) # invert R and BGR -> RGB
return (t,)
class BatchAverageImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply"
CATEGORY = "image/filters"
def apply(self, images):
t = images.detach().clone()
return (torch.mean(t, dim=0, keepdim=True),)
NODE_CLASS_MAPPINGS = {
"AlphaClean": AlphaClean,
"AlphaMatte": AlphaMatte,
@@ -867,11 +990,14 @@ NODE_CLASS_MAPPINGS = {
"BatchNormalizeImage": BatchNormalizeImage,
"DifferenceChecker": DifferenceChecker,
"ImageConstant": ImageConstant,
"ImageConstantHSV": ImageConstantHSV,
"OffsetLatentImage": OffsetLatentImage,
"LatentStats": LatentStats,
"Tonemap": Tonemap,
"UnTonemap": UnTonemap,
"ExposureAdjust": ExposureAdjust,
"ConvertNormals": ConvertNormals,
"BatchAverageImage": BatchAverageImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -890,9 +1016,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"BatchNormalizeImage": "Batch Normalize (Image)",
"DifferenceChecker": "Difference Checker",
"ImageConstant": "Image Constant Color (RGB)",
"ImageConstantHSV": "Image Constant Color (HSV)",
"OffsetLatentImage": "Offset Latent Image",
"LatentStats": "Latent Stats",
"Tonemap": "Tonemap",
"UnTonemap": "UnTonemap",
"ExposureAdjust": "Exposure Adjust",
"ConvertNormals": "Convert Normals",
"BatchAverageImage": "Batch Average Image",
}