Add files via upload
current snapshot - still buggy
This commit is contained in:
Blonbar
@@ -0,0 +1,238 @@
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms as T
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from typing import Any, List
|
||||
|
||||
###############################################################################
|
||||
# Optional: fix relative paths for "data/vgg19_gray.pth" if used
|
||||
###############################################################################
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
orig_torch_load = torch.load
|
||||
def custom_torch_load(path, *args, **kwargs):
|
||||
if path == "data/vgg19_gray.pth":
|
||||
path = os.path.join(current_dir, "data", "vgg19_gray.pth")
|
||||
return orig_torch_load(path, *args, **kwargs)
|
||||
torch.load = custom_torch_load
|
||||
###############################################################################
|
||||
|
||||
# Import from your local code
|
||||
import importlib.abc
|
||||
import importlib.machinery
|
||||
|
||||
class CustomImportFinder(importlib.abc.MetaPathFinder):
|
||||
def __init__(self, base_path):
|
||||
self.base_path = base_path
|
||||
def find_spec(self, fullname, path, target=None):
|
||||
if fullname.startswith('utils.') or fullname == 'utils':
|
||||
parts = fullname.split('.')
|
||||
rel_path = os.path.join(*parts) + ".py"
|
||||
full_path = os.path.join(self.base_path, rel_path)
|
||||
if os.path.exists(full_path):
|
||||
return importlib.machinery.ModuleSpec(
|
||||
name=fullname,
|
||||
loader=importlib.machinery.SourceFileLoader(fullname, full_path),
|
||||
origin=full_path
|
||||
)
|
||||
return None
|
||||
|
||||
sys.meta_path.insert(0, CustomImportFinder(current_dir))
|
||||
|
||||
from models.NonlocalNet import WarpNet, VGG19_pytorch
|
||||
from models.ColorVidNet import ColorVidNet
|
||||
from models.FrameColor import frame_colorization
|
||||
from utils.util import batch_lab2rgb_transpose_mc, tensor_lab2rgb, uncenter_l
|
||||
from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor
|
||||
|
||||
import cv2
|
||||
try:
|
||||
_ = cv2.ximgproc
|
||||
except AttributeError:
|
||||
print("[DeepExemplarImageColorNode] Warning: ximgproc not found. WLS filter may be unavailable.")
|
||||
|
||||
# ComfyUI node registry
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
|
||||
# Paths to your model checkpoints
|
||||
ROOT_DIR = os.path.dirname(__file__)
|
||||
NONLOCAL_TEST_PATH = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth")
|
||||
COLOR_TEST_PATH = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth")
|
||||
VGG_PATH = os.path.join(ROOT_DIR, "data", "vgg19_conv.pth")
|
||||
|
||||
# Global references
|
||||
MODELS_LOADED = False
|
||||
NONLOCAL_NET = None
|
||||
COLOR_NET = None
|
||||
VGG_NET = None
|
||||
|
||||
def load_models_if_needed():
|
||||
global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET
|
||||
if MODELS_LOADED:
|
||||
return
|
||||
print("[DeepExemplarImageColorNode] Loading models...")
|
||||
nonlocal_net = WarpNet(1)
|
||||
color_net = ColorVidNet(7)
|
||||
vgg_net = VGG19_pytorch()
|
||||
|
||||
nonlocal_net.load_state_dict(torch.load(NONLOCAL_TEST_PATH, map_location="cuda"))
|
||||
color_net.load_state_dict(torch.load(COLOR_TEST_PATH, map_location="cuda"))
|
||||
vgg_net.load_state_dict(torch.load(VGG_PATH, map_location="cuda"))
|
||||
|
||||
nonlocal_net.cuda().eval()
|
||||
color_net.cuda().eval()
|
||||
vgg_net.cuda().eval()
|
||||
|
||||
MODELS_LOADED = True
|
||||
NONLOCAL_NET = nonlocal_net
|
||||
COLOR_NET = color_net
|
||||
VGG_NET = vgg_net
|
||||
print("[DeepExemplarImageColorNode] Models loaded.")
|
||||
|
||||
###############################################################################
|
||||
# Utilities to convert Torch Tensors <-> PIL
|
||||
###############################################################################
|
||||
def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
|
||||
"""
|
||||
Expects shape [C,H,W]. Converts to PIL (RGB).
|
||||
Clamps to [0..1], then scales to [0..255].
|
||||
"""
|
||||
image_tensor = image_tensor.detach().cpu().clamp(0,1)
|
||||
np_img = (image_tensor.numpy() * 255).astype(np.uint8)
|
||||
np_img = np.transpose(np_img, (1,2,0)) # [H,W,C]
|
||||
pil_img = Image.fromarray(np_img, mode="RGB")
|
||||
return pil_img
|
||||
|
||||
def pil_to_tensor(pil_img: Image.Image) -> torch.Tensor:
|
||||
"""
|
||||
Converts a PIL (RGB) to Torch float tensor [C,H,W], scaled 0..1.
|
||||
"""
|
||||
return T.ToTensor()(pil_img)
|
||||
|
||||
###############################################################################
|
||||
# Transforms from test.py
|
||||
###############################################################################
|
||||
class ComposeTransforms:
|
||||
def __init__(self, steps):
|
||||
self.steps = steps
|
||||
def __call__(self, img):
|
||||
for step in self.steps:
|
||||
img = step(img)
|
||||
return img
|
||||
|
||||
def build_transform(image_size=(432,768)):
|
||||
return ComposeTransforms([
|
||||
CenterPad(image_size),
|
||||
T.CenterCrop(image_size),
|
||||
RGB2Lab(),
|
||||
ToTensor(),
|
||||
Normalize(),
|
||||
])
|
||||
|
||||
###############################################################################
|
||||
# The Node
|
||||
###############################################################################
|
||||
class DeepExemplarImageColorNode:
|
||||
"""
|
||||
Accepts a single image (Torch Tensor or PIL) + optional boolean WLS, etc.
|
||||
Outputs a single colorized image (Torch Tensor).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"input_image": ("IMAGE",),
|
||||
"wls_filter_on": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"forceInput": False,
|
||||
"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_image"
|
||||
CATEGORY = "DeepExemplar"
|
||||
|
||||
def colorize_image(self,
|
||||
input_image: Any,
|
||||
wls_filter_on: bool,
|
||||
lambda_value: float,
|
||||
sigma_color: float) -> (torch.Tensor,):
|
||||
|
||||
load_models_if_needed()
|
||||
|
||||
# 1) Convert input to PIL
|
||||
pil_img = None
|
||||
if isinstance(input_image, torch.Tensor):
|
||||
# shape [C,H,W]
|
||||
pil_img = tensor_to_pil(input_image)
|
||||
elif isinstance(input_image, Image.Image):
|
||||
pil_img = input_image
|
||||
else:
|
||||
print("[DeepExemplarImageColorNode] Unsupported image type. Returning empty.")
|
||||
return (torch.zeros((3,64,64)),)
|
||||
|
||||
# 2) Prepare transforms
|
||||
transform = build_transform((432,768))
|
||||
input_lab = transform(pil_img).unsqueeze(0).cuda() # shape [1,3,H,W]
|
||||
input_small = F.interpolate(input_lab, scale_factor=0.5, mode="bilinear")
|
||||
|
||||
# We'll treat the reference as itself for a single image (no external ref).
|
||||
ref_tensor_small = input_small.clone()
|
||||
|
||||
# Precompute VGG features
|
||||
with torch.no_grad():
|
||||
I_l = input_small[:,0:1,:,:]
|
||||
I_ab = input_small[:,1:3,:,:]
|
||||
I_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_l), I_ab), dim=1))
|
||||
features_B = VGG_NET(I_rgb, ["r12","r22","r32","r42","r52"], preprocess=True)
|
||||
|
||||
# colorize
|
||||
with torch.no_grad():
|
||||
I_current_ab_predict, _, _ = frame_colorization(
|
||||
I_current_lab=input_small,
|
||||
I_reference_lab=ref_tensor_small,
|
||||
I_last_lab=torch.zeros_like(input_small),
|
||||
features_B=features_B,
|
||||
vggnet=VGG_NET,
|
||||
nonlocal_net=NONLOCAL_NET,
|
||||
color_net=COLOR_NET,
|
||||
feature_noise=0,
|
||||
temperature=1e-10,
|
||||
)
|
||||
|
||||
# Upscale ab
|
||||
ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear") * 1.25
|
||||
|
||||
# optional WLS
|
||||
if wls_filter_on:
|
||||
guide_image = uncenter_l(input_lab[:,0:1,:,:]) * 255/100.0
|
||||
guide_numpy = guide_image[0,0].detach().cpu().numpy().astype(np.uint8)
|
||||
wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter(
|
||||
guide_numpy, lambda_value, sigma_color
|
||||
)
|
||||
a_channel = ab_up[0,0].detach().cpu().numpy()
|
||||
b_channel = ab_up[0,1].detach().cpu().numpy()
|
||||
a_filtered = wls_filter.filter(a_channel)
|
||||
b_filtered = wls_filter.filter(b_channel)
|
||||
a_filt_t = torch.from_numpy(a_filtered).unsqueeze(0).unsqueeze(0)
|
||||
b_filt_t = torch.from_numpy(b_filtered).unsqueeze(0).unsqueeze(0)
|
||||
ab_up = torch.cat([a_filt_t, b_filt_t], dim=1).to(ab_up.device).unsqueeze(0)
|
||||
|
||||
# Convert final to RGB
|
||||
out_rgb = batch_lab2rgb_transpose_mc(input_lab[:,0:1,:,:], ab_up[:1])
|
||||
np_rgb = out_rgb[0].astype(np.uint8)
|
||||
color_pil = Image.fromarray(np_rgb, mode="RGB")
|
||||
|
||||
# Convert back to Torch
|
||||
output_tensor = pil_to_tensor(color_pil) # shape [C,H,W], float 0..1
|
||||
return (output_tensor,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["DeepExemplarImageColorNode"] = DeepExemplarImageColorNode
|
||||
@@ -0,0 +1,365 @@
|
||||
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
|
||||
+14
-1
@@ -1 +1,14 @@
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add the parent directory of 'models' and 'utils' to the Python path
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
if current_dir not in sys.path:
|
||||
sys.path.append(current_dir)
|
||||
|
||||
from .DeepExemplarImageColorNode import NODE_CLASS_MAPPINGS as IMG_MAP
|
||||
from .DeepExemplarVideoColorNode import NODE_CLASS_MAPPINGS as VID_MAP
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_CLASS_MAPPINGS.update(IMG_MAP)
|
||||
NODE_CLASS_MAPPINGS.update(VID_MAP)
|
||||
+1
-1
@@ -3,7 +3,6 @@ tqdm>=4.23.0
|
||||
scipy==1.2
|
||||
torchsummary>=1.3
|
||||
matplotlib>=2.2.2
|
||||
opencv_contrib_python>=4.1.0.25
|
||||
torchvision>=0.3.0
|
||||
scikit_image>=0.15.0
|
||||
torchviz>=0.0.1
|
||||
@@ -17,3 +16,4 @@ future
|
||||
numba
|
||||
pypng
|
||||
opencv-python>=4.5
|
||||
opencv-contrib-python
|
||||
Reference in New Issue
Block a user