189 lines
6.6 KiB
Python
189 lines
6.6 KiB
Python
# © 2025 rndnanthu – Licensed under CC BY-NC 4.0
|
||
import os, torch, numpy as np, cv2, json
|
||
import imageio
|
||
|
||
# Optional EXR support
|
||
try:
|
||
import OpenEXR, Imath
|
||
exr_avail = True
|
||
except ImportError:
|
||
exr_avail = False
|
||
|
||
from .log_transforms import (
|
||
srgb_to_linear,
|
||
rec709_to_linear,
|
||
linear_to_logc3,
|
||
linear_to_clog2,
|
||
linear_to_slog3,
|
||
linear_to_proreslog,
|
||
)
|
||
|
||
BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||
OCIO_FILE = os.path.join(BASE_PATH, "assets", "presets", "ocio_sim_config.json")
|
||
|
||
|
||
class LogReconstructionNode:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
with open(OCIO_FILE, "r") as f:
|
||
cfg = json.load(f)
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"input_space": (cfg["input_spaces"], {"default": cfg["default_input_space"]}),
|
||
"log_transform": (list(cfg["log_transforms"].keys()), {"default": cfg["default_log_transform"]}),
|
||
"use_inverse_tonemap": ("BOOLEAN", {"default": True}),
|
||
"exposure_index": ("FLOAT", {"default": 800.0, "min": 100.0, "max": 6400.0}),
|
||
"auto_analyze": ("BOOLEAN", {"default": True}),
|
||
"local_gain_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
|
||
},
|
||
"optional": {
|
||
"depth_map": ("IMAGE",),
|
||
"normal_map": ("IMAGE",),
|
||
"save_exr": ("BOOLEAN", {"default": False}),
|
||
"exr_filename": ("STRING", {"default": "output.exr"}),
|
||
"save_dpx": ("BOOLEAN", {"default": False}),
|
||
"dpx_filename": ("STRING", {"default": "output.dpx"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "METADATA")
|
||
RETURN_NAMES = ("log_image", "metadata")
|
||
FUNCTION = "run"
|
||
CATEGORY = "rndnanthu/🎨Color Tools"
|
||
|
||
def analyze_image(self, img):
|
||
return {
|
||
"mean": float(np.mean(img)),
|
||
"std": float(np.std(img)),
|
||
"histogram": np.histogram(img, bins=64, range=(0, 1))[0].tolist()
|
||
}
|
||
|
||
def apply_inverse_tonemap(self, img):
|
||
img = np.clip(img, 0.0, 1.0)
|
||
return np.power(img, 2.0) * 1.2 # Soft shoulder boost
|
||
|
||
def inject_local_gain(self, img, depth_map, normal_map, strength):
|
||
H, W, _ = img.shape
|
||
gain = np.ones((H, W), dtype=np.float32)
|
||
|
||
if depth_map is not None:
|
||
d = depth_map[0].cpu().numpy()
|
||
if d.ndim == 3 and d.shape[2] == 3:
|
||
d = cv2.cvtColor(d, cv2.COLOR_RGB2GRAY)
|
||
d = cv2.resize(d, (W, H), interpolation=cv2.INTER_LINEAR)
|
||
gain += (1.0 - d) * strength # Near = more boost
|
||
|
||
if normal_map is not None:
|
||
n = normal_map[0].cpu().numpy()
|
||
if n.ndim == 3 and n.shape[2] == 3:
|
||
g = cv2.cvtColor(n, cv2.COLOR_RGB2GRAY)
|
||
else:
|
||
g = n
|
||
g = cv2.resize(g, (W, H), interpolation=cv2.INTER_LINEAR)
|
||
gain += (1.0 - g) * (strength * 0.5)
|
||
|
||
# Apply gain to all 3 channels
|
||
return np.clip(img * gain[..., None], 0, 1)
|
||
|
||
def export_exr(self, img, path):
|
||
if not exr_avail:
|
||
return
|
||
h, w, _ = img.shape
|
||
hdr = OpenEXR.Header(w, h)
|
||
pt = Imath.PixelType(Imath.PixelType.FLOAT)
|
||
hdr['channels'] = {'R': Imath.Channel(pt), 'G': Imath.Channel(pt), 'B': Imath.Channel(pt)}
|
||
ex = OpenEXR.OutputFile(path, hdr)
|
||
r, g, b = [img[..., i].astype(np.float32).tobytes() for i in range(3)]
|
||
ex.writePixels({'R': r, 'G': g, 'B': b})
|
||
ex.close()
|
||
|
||
def run(self, image, input_space, log_transform, use_inverse_tonemap,
|
||
exposure_index, auto_analyze, local_gain_strength,
|
||
depth_map=None, normal_map=None,
|
||
save_exr=False, exr_filename=None, save_dpx=False, dpx_filename=None):
|
||
|
||
cfg = json.load(open(OCIO_FILE))
|
||
|
||
linearize_func_map = {
|
||
"sRGB": srgb_to_linear,
|
||
"Rec.709": rec709_to_linear,
|
||
"ACEScg": lambda x: x
|
||
}
|
||
linearize = linearize_func_map.get(input_space, lambda x: x)
|
||
|
||
log_func_map = {
|
||
"Alexa_LogC3": linear_to_logc3,
|
||
"Canon_Log2": linear_to_clog2,
|
||
"Sony_SLog3": linear_to_slog3,
|
||
"ProRes_Log": linear_to_proreslog,
|
||
}
|
||
logf = log_func_map.get(log_transform, lambda x: x)
|
||
|
||
# Convert tensor to numpy
|
||
img = image[0].cpu().numpy().astype(np.float32)
|
||
img = np.clip(img, 0, 1)
|
||
|
||
# 1. Linearize color space
|
||
lin = linearize(img)
|
||
metadata = {}
|
||
|
||
# 2. Auto exposure using mean luminance
|
||
if auto_analyze:
|
||
stats = self.analyze_image(lin)
|
||
mean = stats["mean"]
|
||
if mean > 0:
|
||
gain = 0.18 / mean
|
||
lin *= gain
|
||
metadata["auto_gain"] = float(gain)
|
||
|
||
# 3. Apply EI scaling
|
||
lin *= (exposure_index / 800.0)
|
||
|
||
# 4. Apply local gain if maps present
|
||
has_depth = depth_map is not None and isinstance(depth_map, (list, tuple, torch.Tensor))
|
||
has_normal = normal_map is not None and isinstance(normal_map, (list, tuple, torch.Tensor))
|
||
if has_depth or has_normal:
|
||
lin = self.inject_local_gain(lin, depth_map, normal_map, local_gain_strength)
|
||
|
||
# 5. Optional shoulder boost (fake inverse tone map)
|
||
if use_inverse_tonemap:
|
||
lin = self.apply_inverse_tonemap(lin)
|
||
|
||
# 6. Apply log encoding (camera simulation)
|
||
log_img = logf(lin)
|
||
|
||
# 7. Clamp result to [0, 1] for output
|
||
log_img = np.clip(log_img, 0.0, 1.0)
|
||
|
||
# 8. Metadata again after log
|
||
if auto_analyze:
|
||
metadata.update(self.analyze_image(log_img))
|
||
|
||
# 9. Optional EXR/DPX export
|
||
if save_exr:
|
||
os.makedirs("output", exist_ok=True)
|
||
self.export_exr(log_img, os.path.join("output", exr_filename))
|
||
if save_dpx:
|
||
os.makedirs("output", exist_ok=True)
|
||
imageio.imwrite(os.path.join("output", dpx_filename or "output.dpx"),
|
||
(log_img * 65535).astype(np.uint16))
|
||
|
||
# 10. Return tensor + metadata
|
||
tensor = torch.from_numpy(log_img).unsqueeze(0).float()
|
||
metadata.update({
|
||
"input_space": input_space,
|
||
"log_transform": log_transform,
|
||
"exposure_index": exposure_index,
|
||
})
|
||
return (tensor, metadata)
|
||
|
||
|
||
# 🔌 ComfyUI Mappings
|
||
NODE_CLASS_MAPPINGS = {
|
||
"LogReconstructionNode": LogReconstructionNode
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"LogReconstructionNode": "🎯 AI → Camera‑Accurate Log Simulator"
|
||
}
|