420 lines
16 KiB
Python
420 lines
16 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
import cv2
|
|
import math
|
|
from typing import Tuple, Optional, List, Dict, Any
|
|
|
|
# --- 基础工具函数 ---
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def tensor2cv2(image_tensor):
|
|
if len(image_tensor.shape) == 4:
|
|
image_tensor = image_tensor[0]
|
|
img_np = image_tensor.detach().cpu().numpy()
|
|
img_np = np.clip(img_np, 0.0, 1.0)
|
|
# 转换 HWC/CHW
|
|
if img_np.ndim == 3 and img_np.shape[-1] in (3, 4):
|
|
img_hwc = img_np
|
|
elif img_np.ndim == 3 and img_np.shape[0] in (3, 4):
|
|
img_hwc = np.transpose(img_np, (1, 2, 0))
|
|
else:
|
|
raise ValueError(f"Unsupported image layout")
|
|
|
|
img_u8 = (img_hwc * 255.0).astype(np.uint8)
|
|
# 处理RGBA转BGR (OpenCV默认处理)
|
|
if img_u8.shape[2] == 4:
|
|
# 如果是RGBA，我们这里只取RGB做追踪，但后续处理要保留Alpha
|
|
pass
|
|
|
|
# 返回RGB转BGR的结果用于cv2处理
|
|
if img_u8.shape[2] == 4:
|
|
return cv2.cvtColor(img_u8, cv2.COLOR_RGBA2BGR)
|
|
else:
|
|
return cv2.cvtColor(img_u8, cv2.COLOR_RGB2BGR)
|
|
|
|
def parse_color(color_str: str) -> Tuple[int, int, int]:
|
|
color_str = color_str.strip().lstrip('#')
|
|
if ',' in color_str:
|
|
parts = color_str.split(',')
|
|
elif ' ' in color_str:
|
|
parts = color_str.split()
|
|
elif len(color_str) == 6:
|
|
return tuple(int(color_str[i:i+2], 16) for i in (0, 2, 4))
|
|
else:
|
|
parts = [color_str]
|
|
|
|
if len(parts) >= 3:
|
|
return (int(parts[0]), int(parts[1]), int(parts[2]))
|
|
return (255, 0, 0) # Default Red
|
|
|
|
# --- 核心算法 ---
|
|
|
|
def apply_motion_blur(image_pil: Image.Image, velocity_x: float, velocity_y: float, strength: float) -> Image.Image:
|
|
"""应用基于速度的方向性运动模糊"""
|
|
if strength <= 0 or (abs(velocity_x) < 1 and abs(velocity_y) < 1):
|
|
return image_pil
|
|
|
|
# 计算模糊核大小和角度
|
|
magnitude = math.sqrt(velocity_x**2 + velocity_y**2)
|
|
kernel_size = int(magnitude * strength)
|
|
|
|
# 限制核大小，避免过大导致崩溃或极度卡顿
|
|
if kernel_size < 2: return image_pil
|
|
if kernel_size > 50: kernel_size = 50
|
|
|
|
# 确保核大小为奇数
|
|
if kernel_size % 2 == 0: kernel_size += 1
|
|
|
|
# 转换PIL到CV2 (保留Alpha通道)
|
|
img_np = np.array(image_pil)
|
|
|
|
# 检查图像尺寸，避免对过小图像应用模糊
|
|
if img_np.shape[0] < kernel_size or img_np.shape[1] < kernel_size:
|
|
return image_pil
|
|
|
|
# 创建方向性模糊核
|
|
kernel = np.zeros((kernel_size, kernel_size), dtype=np.float32)
|
|
center = kernel_size // 2
|
|
|
|
# 归一化向量
|
|
vec_x = velocity_x / (magnitude + 1e-5)
|
|
vec_y = velocity_y / (magnitude + 1e-5)
|
|
|
|
# 画线生成核
|
|
# 线的起点和终点
|
|
p1 = (int(center - vec_x * (kernel_size/2)), int(center - vec_y * (kernel_size/2)))
|
|
p2 = (int(center + vec_x * (kernel_size/2)), int(center + vec_y * (kernel_size/2)))
|
|
|
|
# 确保坐标在有效范围内
|
|
p1 = (max(0, min(kernel_size-1, p1[0])), max(0, min(kernel_size-1, p1[1])))
|
|
p2 = (max(0, min(kernel_size-1, p2[0])), max(0, min(kernel_size-1, p2[1])))
|
|
|
|
cv2.line(kernel, p1, p2, 1.0, 1)
|
|
|
|
# 归一化核
|
|
kernel_sum = np.sum(kernel)
|
|
if kernel_sum > 0:
|
|
kernel /= kernel_sum
|
|
else:
|
|
# 如果核为空，创建一个简单的中心点核
|
|
kernel[center, center] = 1.0
|
|
|
|
# 应用滤波
|
|
try:
|
|
blurred = cv2.filter2D(img_np, -1, kernel)
|
|
return Image.fromarray(blurred)
|
|
except Exception as e:
|
|
# 如果滤波失败，返回原图像
|
|
print(f"Motion blur failed: {e}")
|
|
return image_pil
|
|
|
|
def find_contour_metrics(frame_cv2: np.ndarray, target_color_rgb: Tuple[int, int, int],
|
|
tolerance: int, processing_width: int,
|
|
prev_center: Optional[Tuple[int, int]]) -> Dict[str, Any]:
|
|
"""
|
|
在降采样后的图像上查找：中心点、旋转角度、面积
|
|
"""
|
|
# 1. 降采样 (Downsampling)
|
|
h, w = frame_cv2.shape[:2]
|
|
scale_ratio = 1.0
|
|
work_frame = frame_cv2
|
|
|
|
if processing_width > 0 and w > processing_width:
|
|
scale_ratio = processing_width / w
|
|
new_h = int(h * scale_ratio)
|
|
work_frame = cv2.resize(frame_cv2, (processing_width, new_h), interpolation=cv2.INTER_AREA)
|
|
|
|
# 2. HSV Masking
|
|
hsv = cv2.cvtColor(work_frame, cv2.COLOR_BGR2HSV)
|
|
target_bgr = np.uint8([[target_color_rgb[::-1]]])
|
|
target_hsv = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2HSV)[0][0]
|
|
|
|
th, ts, tv = int(target_hsv[0]), int(target_hsv[1]), int(target_hsv[2])
|
|
|
|
# 动态容差策略
|
|
h_tol = max(5, int(tolerance * 0.5))
|
|
sv_tol = max(40, int(tolerance * 1.5))
|
|
|
|
lower = np.array([max(0, th - h_tol), max(20, ts - sv_tol), max(20, tv - sv_tol)], dtype=np.uint8)
|
|
upper = np.array([min(180, th + h_tol), min(255, ts + sv_tol), min(255, tv + sv_tol)], dtype=np.uint8)
|
|
|
|
# 环绕色相处理
|
|
if lower[0] > upper[0]:
|
|
# 这种情况通常不需要特殊处理，除非跨越0/180，这里简化逻辑，通常直接clip就好
|
|
pass
|
|
|
|
mask = cv2.inRange(hsv, lower, upper)
|
|
|
|
# 形态学去噪
|
|
kernel_small = np.ones((3, 3), np.uint8)
|
|
kernel_large = np.ones((5, 5), np.uint8)
|
|
mask = cv2.erode(mask, kernel_small, iterations=1)
|
|
mask = cv2.dilate(mask, kernel_large, iterations=2)
|
|
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
if not contours:
|
|
return None
|
|
|
|
valid_candidates = []
|
|
|
|
# 将上一帧的中心点映射到缩小后的坐标系
|
|
prev_center_scaled = None
|
|
if prev_center:
|
|
prev_center_scaled = (prev_center[0] * scale_ratio, prev_center[1] * scale_ratio)
|
|
|
|
for cnt in contours:
|
|
area = cv2.contourArea(cnt)
|
|
if area < 10: continue
|
|
|
|
# 最小外接矩形 (Center(x,y), (w,h), angle)
|
|
rect = cv2.minAreaRect(cnt)
|
|
(cx, cy), (rw, rh), angle = rect
|
|
|
|
# 简单的评分机制
|
|
score = area
|
|
if prev_center_scaled:
|
|
dist = math.sqrt((cx - prev_center_scaled[0])**2 + (cy - prev_center_scaled[1])**2)
|
|
# 距离惩罚：距离越远，分数扣得越多
|
|
score = area - (dist * dist * 0.5)
|
|
|
|
valid_candidates.append({
|
|
'center': (cx / scale_ratio, cy / scale_ratio), # 还原回原分辨率
|
|
'angle': angle,
|
|
'area': area / (scale_ratio * scale_ratio), # 还原回原分辨率
|
|
'score': score,
|
|
'rect_size': (rw, rh)
|
|
})
|
|
|
|
if not valid_candidates:
|
|
return None
|
|
|
|
# 选分最高的
|
|
best = max(valid_candidates, key=lambda x: x['score'])
|
|
return best
|
|
|
|
def smooth_value(current, previous, factor):
|
|
"""通用的一阶平滑"""
|
|
if previous is None: return current
|
|
return previous * (1 - factor) + current * factor
|
|
|
|
def interpolate_missing(data_list):
|
|
"""线性插值填补None"""
|
|
n = len(data_list)
|
|
for i in range(n):
|
|
if data_list[i] is None:
|
|
# Find prev
|
|
prev_idx = i - 1
|
|
while prev_idx >= 0 and data_list[prev_idx] is None: prev_idx -= 1
|
|
# Find next
|
|
next_idx = i + 1
|
|
while next_idx < n and data_list[next_idx] is None: next_idx += 1
|
|
|
|
if prev_idx >= 0 and next_idx < n:
|
|
start = data_list[prev_idx]
|
|
end = data_list[next_idx]
|
|
ratio = (i - prev_idx) / (next_idx - prev_idx)
|
|
data_list[i] = start + (end - start) * ratio
|
|
elif prev_idx >= 0:
|
|
data_list[i] = data_list[prev_idx]
|
|
elif next_idx < n:
|
|
data_list[i] = data_list[next_idx]
|
|
else:
|
|
data_list[i] = 0.0 # Default
|
|
return data_list
|
|
|
|
# --- 节点定义 ---
|
|
|
|
class YCImageColorTrackBindPro:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"video_frames": ("IMAGE",),
|
|
"target_color": ("STRING", {"default": "#FF0000", "multiline": False}),
|
|
"color_tolerance": ("INT", {"default": 30, "min": 0, "max": 255}),
|
|
"smoothing_factor": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 0.99, "step": 0.01}),
|
|
"processing_width": ("INT", {"default": 640, "min": 100, "max": 4096, "step": 10, "label": "Tracking Resolution (px)"}),
|
|
|
|
# 位置微调
|
|
"anchor_offset_x": ("INT", {"default": 0}),
|
|
"anchor_offset_y": ("INT", {"default": 0}),
|
|
"base_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0}),
|
|
|
|
# 高级特性开关
|
|
"enable_rotation": ("BOOLEAN", {"default": False}),
|
|
"rotation_offset": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0}),
|
|
|
|
"enable_auto_scale": ("BOOLEAN", {"default": False, "label": "Auto Z-Depth Scale"}),
|
|
|
|
"motion_blur_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("output_frames",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "YC_VideoCutHelper/Image"
|
|
|
|
def process(self, image, video_frames, target_color, color_tolerance, smoothing_factor, processing_width,
|
|
anchor_offset_x, anchor_offset_y, base_scale,
|
|
enable_rotation, rotation_offset, enable_auto_scale, motion_blur_strength, mask=None):
|
|
|
|
# 1. 准备输入
|
|
target_color_rgb = parse_color(target_color)
|
|
if image.shape[0] > 1: image = image[0:1]
|
|
input_pil_base = tensor2pil(image)
|
|
|
|
if mask is not None:
|
|
if mask.shape[0] > 1: mask = mask[0:1]
|
|
mask_np = mask.cpu().numpy()
|
|
if mask_np.ndim == 3: mask_np = mask_np[0]
|
|
# Resize mask to image size if needed
|
|
if mask_np.shape != (input_pil_base.height, input_pil_base.width):
|
|
m_img = Image.fromarray((mask_np * 255).astype(np.uint8)).resize(input_pil_base.size, Image.Resampling.LANCZOS)
|
|
mask_np = np.array(m_img).astype(np.float32) / 255.0
|
|
|
|
# Apply Mask to Alpha
|
|
if input_pil_base.mode != 'RGBA': input_pil_base = input_pil_base.convert('RGBA')
|
|
r, g, b, a = input_pil_base.split()
|
|
a_np = np.array(a).astype(np.float32) / 255.0
|
|
new_a = (a_np * mask_np * 255).astype(np.uint8)
|
|
input_pil_base.putalpha(Image.fromarray(new_a))
|
|
else:
|
|
if input_pil_base.mode != 'RGBA': input_pil_base = input_pil_base.convert('RGBA')
|
|
|
|
num_frames = video_frames.shape[0]
|
|
|
|
# 2. 追踪阶段 (收集原始数据)
|
|
raw_centers = []
|
|
raw_angles = []
|
|
raw_areas = []
|
|
|
|
last_center = None
|
|
|
|
for i in range(num_frames):
|
|
frame = tensor2cv2(video_frames[i:i+1])
|
|
metrics = find_contour_metrics(frame, target_color_rgb, color_tolerance, processing_width, last_center)
|
|
|
|
if metrics:
|
|
raw_centers.append(metrics['center'])
|
|
raw_angles.append(metrics['angle'])
|
|
raw_areas.append(metrics['area'])
|
|
last_center = metrics['center']
|
|
else:
|
|
raw_centers.append(None)
|
|
raw_angles.append(None)
|
|
raw_areas.append(None)
|
|
|
|
# 3. 数据清洗与插值
|
|
# 分离XY以便插值
|
|
xs = [p[0] if p else None for p in raw_centers]
|
|
ys = [p[1] if p else None for p in raw_centers]
|
|
|
|
fixed_xs = interpolate_missing(xs)
|
|
fixed_ys = interpolate_missing(ys)
|
|
fixed_angles = interpolate_missing(raw_angles)
|
|
fixed_areas = interpolate_missing(raw_areas)
|
|
|
|
# 确定基准面积 (取所有有效面积的中位数，避免初始噪点)
|
|
valid_areas = [a for a in fixed_areas if a > 0]
|
|
base_area_ref = np.median(valid_areas) if valid_areas else 100.0
|
|
|
|
# 4. 平滑与生成
|
|
output_tensors = []
|
|
|
|
# 状态变量
|
|
curr_x, curr_y = fixed_xs[0], fixed_ys[0]
|
|
curr_angle = fixed_angles[0]
|
|
curr_area = fixed_areas[0]
|
|
|
|
# 运动模糊用的前一帧位置
|
|
prev_x, prev_y = curr_x, curr_y
|
|
|
|
# 计算图像初始中心
|
|
img_w_orig, img_h_orig = input_pil_base.size
|
|
|
|
for i in range(num_frames):
|
|
# --- 平滑逻辑 ---
|
|
# 动态调整平滑系数：如果有剧烈运动（距离变化大），减少平滑以防滞后
|
|
# 这里简化处理，直接应用 EMA
|
|
alpha = 1.0 - smoothing_factor
|
|
|
|
# 保存当前位置用于运动模糊计算
|
|
prev_x, prev_y = curr_x, curr_y
|
|
|
|
curr_x = curr_x * (1-alpha) + fixed_xs[i] * alpha
|
|
curr_y = curr_y * (1-alpha) + fixed_ys[i] * alpha
|
|
|
|
# 角度处理 (注意 OpenCV minAreaRect 角度可能是 0-90 或 0-180，这里直接平滑，假设不会突变)
|
|
curr_angle = curr_angle * (1-alpha) + fixed_angles[i] * alpha
|
|
|
|
# 面积处理
|
|
curr_area = curr_area * (1-alpha) + fixed_areas[i] * alpha
|
|
|
|
# --- 变换计算 ---
|
|
|
|
# 1. 缩放
|
|
final_scale = base_scale
|
|
if enable_auto_scale and base_area_ref > 0:
|
|
# 面积比的平方根 = 边长比
|
|
scale_mult = math.sqrt(curr_area / base_area_ref)
|
|
# 限制动态缩放范围，防止无限大或消失
|
|
scale_mult = max(0.2, min(5.0, scale_mult))
|
|
final_scale *= scale_mult
|
|
|
|
new_w = int(img_w_orig * final_scale)
|
|
new_h = int(img_h_orig * final_scale)
|
|
if new_w < 1 or new_h < 1: new_w, new_h = 1, 1
|
|
|
|
img_transformed = input_pil_base.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
|
|
# 2. 旋转
|
|
if enable_rotation:
|
|
# 注意：OpenCV的角度通常是顺时针或逆时针，需根据实际情况调整符号
|
|
# rotation_offset 用于修正初始方向
|
|
total_angle = curr_angle + rotation_offset
|
|
img_transformed = img_transformed.rotate(total_angle, expand=True, resample=Image.Resampling.BICUBIC)
|
|
|
|
# 3. 运动模糊
|
|
if motion_blur_strength > 0 and i > 0:
|
|
# 计算速度向量（当前位置 - 前一帧位置）
|
|
vx = curr_x - prev_x
|
|
vy = curr_y - prev_y
|
|
img_transformed = apply_motion_blur(img_transformed, vx, vy, motion_blur_strength)
|
|
|
|
# 4. 合成
|
|
frame_tensor = video_frames[i:i+1]
|
|
frame_pil = tensor2pil(frame_tensor)
|
|
bg_w, bg_h = frame_pil.size
|
|
|
|
# 计算粘贴坐标 (中心对齐)
|
|
paste_w, paste_h = img_transformed.size
|
|
paste_x = int(curr_x - paste_w / 2 + anchor_offset_x)
|
|
paste_y = int(curr_y - paste_h / 2 + anchor_offset_y)
|
|
|
|
# PIL粘贴处理透明度
|
|
comp = frame_pil.copy().convert('RGBA')
|
|
comp.paste(img_transformed, (paste_x, paste_y), img_transformed)
|
|
|
|
output_tensors.append(pil2tensor(comp.convert('RGB')))
|
|
|
|
return (torch.cat(output_tensors, dim=0),)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"YCImageColorTrackBindPro": YCImageColorTrackBindPro,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"YCImageColorTrackBindPro": "Image Color Track Bind (Pro)",
|
|
} |