import os import sys import torch import torch.nn.functional as F import numpy as np from typing import Any, List from tqdm import tqdm import cv2 from PIL import Image ############################################################################### # Global references for the colorization nets ############################################################################### MODELS_LOADED = False NONLOCAL_NET = None COLOR_NET = None VGG_NET = None ############################################################################### # Paths from install.py ############################################################################### current_dir = os.path.dirname(os.path.abspath(__file__)) nonlocal_ckpt = os.path.join(current_dir, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth") color_ckpt = os.path.join(current_dir, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth") vgg_ckpt = os.path.join(current_dir, "data", "vgg19_conv.pth") ############################################################################### # Helper: get height,width from PIL or Torch ############################################################################### def get_hw(image_or_tensor): """Return (height, width) from a PIL image or Torch [C,H,W] or [N,C,H,W].""" if isinstance(image_or_tensor, Image.Image): return image_or_tensor.height, image_or_tensor.width elif isinstance(image_or_tensor, torch.Tensor): # shape [C,H,W] => (H,W) # shape [N,C,H,W] => pick the first if image_or_tensor.dim()==4: if image_or_tensor.size(0)==0: return 0,0 return image_or_tensor.shape[-2], image_or_tensor.shape[-1] elif image_or_tensor.dim()==3: return image_or_tensor.shape[-2], image_or_tensor.shape[-1] else: return 0,0 else: return 0,0 ############################################################################### # CenterPadToRef: Pad or Crop to (ref_h, ref_w) ############################################################################### import numpy as np import torchvision.transforms as Tlib class CenterPadToRef: """ If the image is smaller in a dimension => pad If bigger => center-crop Expects a PIL image as input. """ def __init__(self, ref_h, ref_w): self.ref_h = ref_h self.ref_w = ref_w def __call__(self, pil_img): arr = np.array(pil_img) h_old, w_old = arr.shape[:2] out_h, out_w = self.ref_h, self.ref_w # If smaller => pad, if bigger => crop # We'll do a naive approach if h_old==out_h and w_old==out_w: return pil_img # create an output c_old = 1 if arr.ndim==3 and arr.shape[2]==1 else 3 out_arr = np.zeros((out_h, out_w, c_old), dtype=arr.dtype) # if bigger => crop if h_old>out_h: start_h = (h_old - out_h)//2 arr = arr[start_h:start_h+out_h,:,:] h_old = out_h if w_old>out_w: start_w = (w_old - out_w)//2 arr = arr[:, start_w:start_w+out_w, :] w_old = out_w # if smaller => pad off_h = (out_h - h_old)//2 if out_h>h_old else 0 off_w = (out_w - w_old)//2 if out_w>w_old else 0 h_end = off_h + h_old w_end = off_w + w_old out_arr[off_h:h_end, off_w:w_end] = arr return Image.fromarray(out_arr) ############################################################################### # Convert PIL or Torch to Lab tensor (L => L-50) ############################################################################### from skimage import color def convert_any_image_to_lab_tensor(item): """ If PIL => do np.array => rgb2lab => (L-50) If Torch => shape [C,H,W] or [N,C,H,W], convert to CPU np => rgb2lab => (L-50) """ if isinstance(item, Image.Image): arr = np.array(item).astype(np.float32) if arr.ndim==2: arr = np.stack([arr,arr,arr], axis=-1) arr /= 255.0 lab = color.rgb2lab(arr) lab_t = torch.from_numpy(lab.transpose(2,0,1)) lab_t[0:1,:,:] = lab_t[0:1,:,:]-50.0 return lab_t elif isinstance(item, torch.Tensor): # if shape [N,C,H,W], pick first if item.dim()==4: if item.size(0)==0: return None item = item[0] # now [C,H,W] arr = item.detach().cpu().float() # guess scale if arr.max()>2.0: arr = arr/255.0 np_arr = arr.numpy() np_arr = np_arr.transpose(1,2,0) # [H,W,C] lab = color.rgb2lab(np_arr) lab_t = torch.from_numpy(lab.transpose(2,0,1)) lab_t[0:1,:,:] = lab_t[0:1,:,:]-50.0 return lab_t else: return None ############################################################################### # final step: batch_lab2rgb_transpose_mc, optional WLS ############################################################################### from utils.util import batch_lab2rgb_transpose_mc, tensor_lab2rgb, uncenter_l from models.NonlocalNet import WarpNet, VGG19_pytorch from models.ColorVidNet import ColorVidNet from models.FrameColor import frame_colorization ############################################################################### # The Node ############################################################################### NODE_CLASS_MAPPINGS = {} class DeepExemplarVideoColorNodeTestPyLike: """ Node that: - checks if reference_image is PIL or Torch - obtains (ref_h, ref_w) - optional center pad/crop each frame to (ref_h, ref_w) - convert to Lab => downsample => colorize => upsample => optional WLS => final batch_lab2rgb - frame_propagate => first frame is reference - WLS defaults: lambda=500, sigma=4 """ @classmethod def INPUT_TYPES(cls): return { "required": { "video_frames": ("IMAGE",), "reference_image": ("IMAGE",), "frame_propagate": ("BOOLEAN", { "default": False, "widget": "checkbox" }), "use_center_pad": ("BOOLEAN", { "default": False, "widget": "checkbox" }), "use_center_crop": ("BOOLEAN", { "default": False, "widget": "checkbox" }), "wls_filter_on": ("BOOLEAN", { "default": True, "widget": "checkbox" }), "lambda_value": ("FLOAT", { "default": 500.0, "min": 0, "max": 2000 }), "sigma_color": ("FLOAT", { "default": 4.0, "min": 0, "max": 50 }), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "colorize_video" CATEGORY = "Video/DeepExemplar" def load_models_if_needed(self): global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET if MODELS_LOADED: return print("[DeepExemplarVideoColorNodeTestPyLike] Loading models from local checkpoints...") nonlocal_net = WarpNet(1) color_net = ColorVidNet(7) vgg_net = VGG19_pytorch() nonlocal_net.load_state_dict(torch.load(nonlocal_ckpt, map_location="cuda")) color_net.load_state_dict(torch.load(color_ckpt, map_location="cuda")) vgg_net.load_state_dict(torch.load(vgg_ckpt, map_location="cuda")) nonlocal_net.cuda().eval() color_net.cuda().eval() vgg_net.cuda().eval() NONLOCAL_NET = nonlocal_net COLOR_NET = color_net VGG_NET = vgg_net MODELS_LOADED = True print("[DeepExemplarVideoColorNodeTestPyLike] Models loaded.") def colorize_video(self, video_frames: Any, reference_image: Any, frame_propagate: bool, use_center_pad: bool, use_center_crop: bool, wls_filter_on: bool, lambda_value: float, sigma_color: float ) -> (List[np.ndarray],): self.load_models_if_needed() # if frame_propagate => first frame is reference # unify frames => list frames_list = [] if isinstance(video_frames, list): frames_list = video_frames elif isinstance(video_frames, torch.Tensor): # shape [N,C,H,W] or [C,H,W] if video_frames.dim()==4: frames_list = [video_frames[i] for i in range(video_frames.size(0))] else: frames_list = [video_frames] elif isinstance(video_frames, Image.Image): frames_list = [video_frames] if frame_propagate and len(frames_list)>0: reference_image = frames_list[0] # get (ref_h, ref_w) ref_h, ref_w = get_hw(reference_image) if ref_h<1 or ref_w<1: print("[DeepExemplarVideoColorNodeTestPyLike] reference_image is zero dimension.") return ([],) # build optional pad/crop pad_transform = CenterPadToRef(ref_h, ref_w) if use_center_pad else None import torchvision.transforms as Tlib crop_transform = Tlib.CenterCrop((ref_h, ref_w)) if use_center_crop else None # Convert reference => lab # if it's PIL or Torch # if pad/crop => must be PIL # so if reference_image is Torch, we convert to PIL if needed ref_pil = reference_image if isinstance(reference_image, torch.Tensor): # convert to PIL ref_pil = self.torch_to_pil(reference_image) if pad_transform: ref_pil = pad_transform(ref_pil) if crop_transform: ref_pil = crop_transform(ref_pil) ref_lab_large = convert_any_image_to_lab_tensor(ref_pil).unsqueeze(0).cuda() ref_lab_small = F.interpolate(ref_lab_large, scale_factor=0.5, mode="bilinear", align_corners=False) with torch.no_grad(): ref_l = ref_lab_small[:,0:1,:,:]+50.0 ref_ab = ref_lab_small[:,1:3,:,:] ref_unnorm = torch.cat([ref_l, ref_ab], dim=1) ref_rgb = tensor_lab2rgb(ref_unnorm) features_B = VGG_NET(ref_rgb, ["r12","r22","r32","r42","r52"], preprocess=True) last_lab_predict = ref_lab_small.clone() if frame_propagate else None out_frames = [] for idx, item in enumerate(tqdm(frames_list, desc="[DeepExemplar Node]")): # unify item => PIL if isinstance(item, torch.Tensor): pil_frame = self.torch_to_pil(item) else: pil_frame = item # pad/crop if toggles if pad_transform: pil_frame = pad_transform(pil_frame) if crop_transform: pil_frame = crop_transform(pil_frame) # convert to lab lab_large = convert_any_image_to_lab_tensor(pil_frame) if lab_large is None: print(f"[DeepExemplar] skipping frame {idx}, invalid.") continue lab_large = lab_large.unsqueeze(0).cuda() lab_small = F.interpolate(lab_large, scale_factor=0.5, mode="bilinear", align_corners=False) IA_l = lab_small[:,0:1,:,:] if last_lab_predict is None: last_lab_predict = torch.zeros_like(lab_small) with torch.no_grad(): I_current_ab_predict, _, _ = frame_colorization( lab_small, ref_lab_small, last_lab_predict, features_B, VGG_NET, NONLOCAL_NET, COLOR_NET, False, 1e-10 ) last_lab_predict = torch.cat([IA_l, I_current_ab_predict], dim=1) ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear", align_corners=False)*1.25 if wls_filter_on: guide_l = (lab_large[:,0:1,:,:]+50.0).clamp(0,100)*255.0/100.0 guide_np = guide_l[0,0].detach().cpu().numpy().astype(np.uint8) wlsf = cv2.ximgproc.createFastGlobalSmootherFilter(guide_np, lambda_value, sigma_color) a_np = ab_up[0,0].detach().cpu().numpy() b_np = ab_up[0,1].detach().cpu().numpy() a_filt = wlsf.filter(a_np) b_filt = wlsf.filter(b_np) a_t = torch.from_numpy(a_filt).unsqueeze(0).unsqueeze(0).cuda() b_t = torch.from_numpy(b_filt).unsqueeze(0).unsqueeze(0).cuda() ab_up = torch.cat([a_t,b_t], dim=1) out_rgb = batch_lab2rgb_transpose_mc(lab_large[:,0:1,:,:]+50.0, ab_up) np_rgb = out_rgb[0].astype(np.uint8) out_frames.append(np_rgb) return (out_frames,) def torch_to_pil(self, tensor: torch.Tensor): """ Convert a Torch [C,H,W] or [N,C,H,W] in 0..1 or 0..255 to a PIL image. We'll just pick the first if [N,C,H,W]. """ if tensor.dim()==4: if tensor.size(0)==0: # return blank arr = np.zeros((1,1,3), dtype=np.uint8) return Image.fromarray(arr) tensor = tensor[0] arr = tensor.detach().cpu().float() if arr.max()>2.0: arr = arr/255.0 arr_np = arr.numpy() # shape [C,H,W] arr_np = arr_np.transpose(1,2,0) # [H,W,C] arr_np = (arr_np*255).clip(0,255).astype(np.uint8) pil_img = Image.fromarray(arr_np) return pil_img NODE_CLASS_MAPPINGS["DeepExemplarVideoColorNodeTestPyLike"] = DeepExemplarVideoColorNodeTestPyLike