91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
import torch
|
|
import time
|
|
from PIL import Image, ImageEnhance
|
|
from .imagefunc import log, tensor2pil, pil2tensor
|
|
from .imagefunc import gamma_trans, depthblur_image, radialblur_image, vignette_image, filmgrain_image
|
|
|
|
|
|
|
|
class Film:
|
|
|
|
def __init__(self):
|
|
self.NODE_NAME = 'Film'
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ), #
|
|
"center_x": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
|
|
"center_y": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
|
|
"saturation": ("FLOAT", {"default": 1, "min": 0.01, "max": 3, "step": 0.01}),
|
|
"vignette_intensity": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
|
|
"grain_power": ("FLOAT", {"default": 0.15, "min": 0, "max": 1, "step": 0.01}),
|
|
"grain_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10, "step": 0.1}),
|
|
"grain_sat": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
|
|
"grain_shadows": ("FLOAT", {"default": 0.6, "min": 0, "max": 1, "step": 0.01}),
|
|
"grain_highs": ("FLOAT", {"default": 0.2, "min": 0, "max": 1, "step": 0.01}),
|
|
"blur_strength": ("INT", {"default": 90, "min": 0, "max": 256, "step": 1}),
|
|
"blur_focus_spread": ("FLOAT", {"default": 2.2, "min": 0.1, "max": 8, "step": 0.1}),
|
|
"focal_depth": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"depth_map": ("IMAGE",), #
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = 'film'
|
|
CATEGORY = '😺dzNodes/LayerFilter'
|
|
|
|
def film(self, image, center_x, center_y, saturation, vignette_intensity,
|
|
grain_power, grain_scale, grain_sat, grain_shadows, grain_highs,
|
|
blur_strength, blur_focus_spread, focal_depth,
|
|
depth_map=None
|
|
):
|
|
|
|
ret_images = []
|
|
seed = int(time.time())
|
|
for i in image:
|
|
i = torch.unsqueeze(i, 0)
|
|
_canvas = tensor2pil(i).convert('RGB')
|
|
|
|
if saturation != 1:
|
|
color_image = ImageEnhance.Color(_canvas)
|
|
_canvas = color_image.enhance(factor= saturation)
|
|
|
|
if blur_strength:
|
|
if depth_map is not None:
|
|
depth_map = tensor2pil(depth_map).convert('L').convert('RGB')
|
|
if depth_map.size != _canvas.size:
|
|
depth_map.resize((_canvas.size), Image.BILINEAR)
|
|
_canvas = depthblur_image(_canvas, depth_map, blur_strength, focal_depth, blur_focus_spread)
|
|
else:
|
|
_canvas = radialblur_image(_canvas, blur_strength, center_x, center_y, blur_focus_spread * 2)
|
|
|
|
if vignette_intensity:
|
|
# adjust image gamma and saturation
|
|
_canvas = gamma_trans(_canvas, 1 - vignette_intensity / 3)
|
|
color_image = ImageEnhance.Color(_canvas)
|
|
_canvas = color_image.enhance(factor= 1+ vignette_intensity / 3)
|
|
# add vignette
|
|
_canvas = vignette_image(_canvas, vignette_intensity, center_x, center_y)
|
|
|
|
if grain_power:
|
|
_canvas = filmgrain_image(_canvas, grain_scale, grain_power, grain_shadows, grain_highs, grain_sat, seed=seed)
|
|
seed += 1
|
|
ret_image = _canvas
|
|
ret_images.append(pil2tensor(ret_image))
|
|
|
|
log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish')
|
|
return (torch.cat(ret_images, dim=0),)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerFilter: Film": Film
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerFilter: Film": "LayerFilter: Film"
|
|
} |