spike temporalnet preprocess

2025-05-01 12:09:22 -04:00
parent 5703bb78ca
commit 4cadd5ca12
1 changed files with 198 additions and 1 deletions
@@ -1,9 +1,12 @@
 import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math

 from ....src.realtimenodes.control_base import ControlNodeBase
 from ....src.utils.general import AlwaysEqualProxy

-
 class StateResetNode(ControlNodeBase):
    """Node that resets all control node states when triggered"""

@@ -157,4 +160,198 @@ class SetStateNode(ControlNodeBase):

        return (value,)

+def gaussian_blur_2d(x, kernel_size=9, sigma=1.0):
+    """
+    Apply 2D Gaussian blur to input tensor.
+    
+    Args:
+        x (torch.Tensor): Input tensor in BCHW format
+        kernel_size (int): Size of the Gaussian kernel
+        sigma (float): Standard deviation of the Gaussian kernel
+    """
+    # Create 1D Gaussian kernel
+    kernel_1d = torch.linspace(-(kernel_size // 2), kernel_size // 2, kernel_size)
+    kernel_1d = torch.exp(-0.5 * (kernel_1d / sigma) ** 2)
+    kernel_1d = kernel_1d / kernel_1d.sum()
+    
+    # Create 2D Gaussian kernel
+    kernel_2d = torch.outer(kernel_1d, kernel_1d)
+    kernel_2d = kernel_2d / kernel_2d.sum()
+    
+    # Reshape kernel for convolution
+    kernel_2d = kernel_2d.view(1, 1, kernel_size, kernel_size)
+    kernel_2d = kernel_2d.repeat(x.shape[1], 1, 1, 1)
+    
+    # Move kernel to same device as input
+    kernel_2d = kernel_2d.to(x.device)
+    
+    # Apply padding
+    padding = kernel_size // 2
+    
+    # Apply convolution
+    return F.conv2d(x, kernel_2d, padding=padding, groups=x.shape[1])
+
+class RealTimeFlowNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Lightweight CNN for real-time flow estimation
+        self.conv1 = nn.Conv2d(2, 16, 3, padding=1)
+        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
+        self.conv3 = nn.Conv2d(32, 2, 3, padding=1)
+        
+        # Initialize weights for better flow estimation
+        nn.init.kaiming_normal_(self.conv1.weight, mode='fan_out', nonlinearity='relu')
+        nn.init.kaiming_normal_(self.conv2.weight, mode='fan_out', nonlinearity='relu')
+        nn.init.kaiming_normal_(self.conv3.weight, mode='fan_out', nonlinearity='relu')
+        
+    def forward(self, x):
+        # x should be in BCHW format
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = self.conv3(x)
+        return x
+
+class TemporalNetV2Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "current_frame": ("IMAGE",),
+                "previous_frame": ("IMAGE",),
+                "flow_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
+                "blur_sigma": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+            }
+        }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "preprocess"
+    CATEGORY = "controlnet/preprocessors"
+
+    def __init__(self):
+        self.flow_net = RealTimeFlowNet()
+        self.flow_net.eval()
+        
+    def preprocess(self, current_frame, previous_frame, flow_scale=1.0, blur_sigma=1.0):
+        # Ensure inputs are on the same device
+        device = current_frame.device
+        self.flow_net = self.flow_net.to(device)
+        
+        # Convert to grayscale and ensure proper format
+        # Input is BHWC, convert to BCHW for CNN
+        current_gray = torch.mean(current_frame, dim=-1, keepdim=True).movedim(-1, 1)
+        prev_gray = torch.mean(previous_frame, dim=-1, keepdim=True).movedim(-1, 1)
+        
+        # Stack frames for flow estimation
+        stacked = torch.cat([prev_gray, current_gray], dim=1)
+        
+        # Estimate flow
+        with torch.no_grad():
+            flow = self.flow_net(stacked)
+            
+        # Scale the flow
+        flow = flow * flow_scale
+        
+        # Apply Gaussian blur if needed
+        if blur_sigma > 0:
+            flow = gaussian_blur_2d(flow, kernel_size=9, sigma=blur_sigma)
+        
+        # Visualize flow as RGB (similar to flow_to_image in torchvision)
+        # This converts 2-channel flow to 3-channel RGB visualization
+        flow_rgb = self.flow_to_rgb(flow)
+        
+        # Convert flow back to BHWC format
+        flow_rgb = flow_rgb.movedim(1, -1)
+        
+        # Convert previous frame to correct format if needed (ensure BHWC)
+        if len(previous_frame.shape) == 3:
+            previous_frame = previous_frame.unsqueeze(0)
+        
+        # Combine previous frame RGB and flow RGB into 6-channel tensor (matching reference)
+        combined = torch.cat([
+            previous_frame,  # Previous frame RGB (3 channels)
+            flow_rgb,        # Flow visualization (3 channels)
+        ], dim=-1)
+        
+        return (combined,)
+    
+    def flow_to_rgb(self, flow):
+        """
+        Convert optical flow to RGB visualization similar to torchvision's flow_to_image
+        
+        Args:
+            flow: optical flow tensor of shape [B, 2, H, W]
+        Returns:
+            RGB visualization tensor of shape [B, 3, H, W]
+        """
+        B, _, H, W = flow.shape
+        
+        # Calculate flow magnitude and angle
+        flow_x = flow[:, 0]
+        flow_y = flow[:, 1]
+        magnitude = torch.sqrt(flow_x**2 + flow_y**2)
+        angle = torch.atan2(flow_y, flow_x)
+        
+        # Normalize magnitude for better visualization
+        max_mag = torch.max(magnitude.view(B, -1), dim=1)[0].view(B, 1, 1)
+        max_mag = torch.clamp(max_mag, min=1e-4)
+        magnitude = torch.clamp(magnitude / max_mag, 0, 1)
+        
+        # Convert angle and magnitude to RGB using HSV->RGB conversion
+        # Hue = angle, Saturation = 1, Value = magnitude
+        angle_normalized = (angle / (2 * math.pi) + 0.5) % 1.0
+        
+        # HSV to RGB conversion
+        h = angle_normalized * 6
+        i = torch.floor(h)
+        f = h - i
+        p = torch.zeros_like(magnitude)
+        q = 1 - f
+        t = f
+        
+        # Initialize RGB channels
+        r = torch.zeros_like(magnitude)
+        g = torch.zeros_like(magnitude)
+        b = torch.zeros_like(magnitude)
+        
+        # Case 0: h in [0,1)
+        mask = (i == 0)
+        r[mask] = magnitude[mask]
+        g[mask] = magnitude[mask] * t[mask]
+        b[mask] = p[mask]
+        
+        # Case 1: h in [1,2)
+        mask = (i == 1)
+        r[mask] = magnitude[mask] * q[mask]
+        g[mask] = magnitude[mask]
+        b[mask] = p[mask]
+        
+        # Case 2: h in [2,3)
+        mask = (i == 2)
+        r[mask] = p[mask]
+        g[mask] = magnitude[mask]
+        b[mask] = magnitude[mask] * t[mask]
+        
+        # Case 3: h in [3,4)
+        mask = (i == 3)
+        r[mask] = p[mask]
+        g[mask] = magnitude[mask] * q[mask]
+        b[mask] = magnitude[mask]
+        
+        # Case 4: h in [4,5)
+        mask = (i == 4)
+        r[mask] = magnitude[mask] * t[mask]
+        g[mask] = p[mask]
+        b[mask] = magnitude[mask]
+        
+        # Case 5: h in [5,6)
+        mask = (i == 5)
+        r[mask] = magnitude[mask]
+        g[mask] = p[mask]
+        b[mask] = magnitude[mask] * q[mask]
+        
+        # Stack RGB channels
+        rgb = torch.stack([r, g, b], dim=1)
+        
+        return rgb
+