66 lines
1.7 KiB
Python
66 lines
1.7 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import numpy as np
|
|
|
|
|
|
class ColorModPivot:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"pivot": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"move": ("FLOAT", {"default": 0.0, "min": -2.0, "max": 2.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "mod_pivot"
|
|
CATEGORY = "image/postprocessing"
|
|
TITLE = "ColorMod (move pivot)"
|
|
|
|
def mod_pivot(self, image, pivot, move):
|
|
pivot_map = torch.ones(image.shape) * pivot
|
|
image_high = torch.maximum(image, pivot_map) - pivot
|
|
image_low = torch.minimum(image, pivot_map)
|
|
|
|
image_high = image_high * (1/(1-pivot)) * (1-(pivot + move))
|
|
image_low = image_low * (1/pivot) * (pivot + move)
|
|
out = torch.clip((image_high + image_low), 0.0, 1.0)
|
|
return (out,)
|
|
|
|
|
|
class ColorModEdges:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"low": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"pivot": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"high": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "mod_edges"
|
|
CATEGORY = "image/postprocessing"
|
|
TITLE = "ColorMod (edges)"
|
|
|
|
def mod_edges(self, image, low, pivot, high):
|
|
pivot_map = torch.ones(image.shape) * pivot
|
|
image_high = torch.maximum(image, pivot_map) - pivot
|
|
image_low = torch.minimum(image, pivot_map)
|
|
|
|
image_low = image_low * low + pivot * (1-low)
|
|
image_high = image_high * high
|
|
out = torch.clip((image_high + image_low), 0.0, 1.0)
|
|
return (out,)
|