From bdde8605e15917b1bc423b869f6510a9010ba949 Mon Sep 17 00:00:00 2001 From: Blonbar <117039186+jonstreeter@users.noreply.github.com> Date: Sun, 16 Mar 2025 09:09:28 -0500 Subject: [PATCH] Delete DeepExemplarVideoColorNode.py outdated --- DeepExemplarVideoColorNode.py | 365 ---------------------------------- 1 file changed, 365 deletions(-) delete mode 100644 DeepExemplarVideoColorNode.py diff --git a/DeepExemplarVideoColorNode.py b/DeepExemplarVideoColorNode.py deleted file mode 100644 index b9a2192..0000000 --- a/DeepExemplarVideoColorNode.py +++ /dev/null @@ -1,365 +0,0 @@ -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