nodes/mcDandy-more_math
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

motion functions and permute fix

Browse Source
This commit is contained in:
mcDandy
2026-02-12 08:51:20 +01:00
parent 0ba8c8cf26
commit 3d9e37b602
12 changed files with 2312 additions and 1543 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+19
View File
@@ -170,6 +170,24 @@ You can also get the node from comfy manager under the name of More math.
- `matmul(a, b)`: Matrix multiplication. For 1D vectors, performs dot product. For 2D+ tensors, performs standard matrix multiplication following NumPy rules.
- `cross(a, b)`: Computes the cross product (vector product) of two 3D vectors. Both inputs must have last dimension = 3. Returns a vector perpendicular to both inputs.
### Optical Flow
- `rife(img1, img2, [tiling_size, iterations, multi_scale])`: Calculates optical flow from `img1` to `img2` using RAFT.
- `img1`, `img2`: Images or video frames [B, H, W, C].
- `tiling_size`: (Optional) Size of tiles for high-res images. Default is `0` (auto-tile if > 2MPx). Supports fractions (e.g., `0.5`).
- `iterations`: (Optional) Number of flow updates (default: 12).
- `multi_scale`: (Optional) Whether to use a global pass for large movements (default: false (0)).
- **Returns**: Flow vectors [B, H, W, 2].
- `motion_mask(flow)`: Generates an occlusion/motion mask from optical flow vectors.
- `flow`: Flow vectors [B, H, W, 2].
- **Returns**: Mask [B, H, W] in range [0, 1].
- `flow_to_image(flow)` or `flow_view(flow)`: Converts flow vectors to an RGB image for visualization.
- `flow`: Flow vectors [B, H, W, 2].
- **Returns**: RGB image [B, H, W, 3].
- `flow_apply(image, flow)` or `apply_flow(image, flow)`: Warps an image using optical flow vectors.
- `image`: Image [B, H, W, C].
- `flow`: Flow vectors [B, H, W, 2] from `rife()`.
- **Returns**: Warped image [B, H, W, C].
### FFT (Tensor Only)
@@ -191,6 +209,7 @@ You can also get the node from comfy manager under the name of More math.
### Random Distributions
Generates random noise with default shape of aither first input or maximum of input sizes, depending on node setting.
- `random_normal(seed,[shape])` or `randn` or `noise`: generates a random tensor with normal distribution (var=1, mean=0).
- `random_uniform(seed,[shape])` or `rand`: generates a random tensor with uniform distribution [0, 1).
- `random_exponential(seed, lambda,[shape])` or `rande`: generates a random tensor with exponential distribution.
more_math/Parser/MathExpr.g4
+14 -4
View File
@@ -144,10 +144,12 @@ func1:
| ARGSORT LPAREN expr (COMMA expr)? RPAREN # ArgsortFunc
| ARGMIN LPAREN expr RPAREN # ArgminFunc
| ARGMAX LPAREN expr RPAREN # ArgmaxFunc
| SOFTMAX LPAREN expr RPAREN # SoftmaxFunc
| SOFTMIN LPAREN expr RPAREN # SoftminFunc
| SOFTMAX LPAREN expr RPAREN # SoftmaxFunc
| SOFTMIN LPAREN expr RPAREN # SoftminFunc
| UNIQUE LPAREN expr RPAREN # UniqueFunc
| FLATTEN LPAREN expr RPAREN # FlattenFunc;
| FLATTEN LPAREN expr RPAREN # FlattenFunc
| MOTION_MASK LPAREN expr RPAREN # MotionMaskFunc
| FLOW_TO_IMAGE LPAREN expr RPAREN # FlowToImageFunc;
func2:
POWE LPAREN expr COMMA expr RPAREN # PowFunc
@@ -174,7 +176,11 @@ func2:
| TENSOR LPAREN indexExpr (COMMA expr)? RPAREN # EmptyTensorFunc
| PAD LPAREN expr COMMA expr RPAREN # PadFunc
| CROSS LPAREN expr COMMA expr RPAREN # CrossFunc
| MATMUL LPAREN expr COMMA expr RPAREN # MatmulFunc;
| MATMUL LPAREN expr COMMA expr RPAREN # MatmulFunc
| FLOW_APPLY LPAREN expr COMMA expr RPAREN # FlowApplyFunc
| RIFE LPAREN expr COMMA expr (
COMMA expr (COMMA expr (COMMA expr)?)?
)? RPAREN # RifeFunc;
func3:
CLAMP LPAREN expr COMMA expr COMMA expr RPAREN # ClampFunc
@@ -343,6 +349,10 @@ TIMESTAMP: 'timestamp' | 'now';
NONE: 'None' | 'none' | 'NULL' | 'null';
BREAK: 'break';
CONTINUE: 'continue';
RIFE: 'rife';
MOTION_MASK: 'motion_mask';
FLOW_TO_IMAGE: 'flow_to_image' | 'flow_view';
FLOW_APPLY: 'flow_apply' | 'apply_flow';
TENSOR: 'tensor';
more_math/Parser/MathExpr.interp
+9 -1
View File
File diff suppressed because one or more lines are too long
more_math/Parser/MathExpr.tokens
+64 -58
View File
@@ -119,38 +119,42 @@ TIMESTAMP=118
NONE=119
BREAK=120
CONTINUE=121
TENSOR=122
PLUS=123
MINUS=124
MULT=125
DIV=126
MOD=127
POW=128
GE=129
GT=130
LE=131
LT=132
EQ=133
EQUEALS=134
NE=135
PIPE=136
LPAREN=137
RPAREN=138
COMMA=139
SEMICOLON=140
ARROW=141
LBRACKET=142
RBRACKET=143
QUESTION=144
COLON=145
LBRACE=146
RBRACE=147
NUMBER=148
CONSTANT=149
VARIABLE=150
SL_COMMENT=151
ML_COMMENT=152
WS=153
RIFE=122
MOTION_MASK=123
FLOW_TO_IMAGE=124
FLOW_APPLY=125
TENSOR=126
PLUS=127
MINUS=128
MULT=129
DIV=130
MOD=131
POW=132
GE=133
GT=134
LE=135
LT=136
EQ=137
EQUEALS=138
NE=139
PIPE=140
LPAREN=141
RPAREN=142
COMMA=143
SEMICOLON=144
ARROW=145
LBRACKET=146
RBRACKET=147
QUESTION=148
COLON=149
LBRACE=150
RBRACE=151
NUMBER=152
CONSTANT=153
VARIABLE=154
SL_COMMENT=155
ML_COMMENT=156
WS=157
'sin'=1
'cos'=2
'tan'=3
@@ -242,29 +246,31 @@ WS=153
'flatten'=115
'break'=120
'continue'=121
'tensor'=122
'+'=123
'-'=124
'*'=125
'/'=126
'%'=127
'^'=128
'>='=129
'>'=130
'<='=131
'<'=132
'=='=133
'='=134
'!='=135
'|'=136
'('=137
')'=138
','=139
';'=140
'->'=141
'['=142
']'=143
'?'=144
':'=145
'{'=146
'}'=147
'rife'=122
'motion_mask'=123
'tensor'=126
'+'=127
'-'=128
'*'=129
'/'=130
'%'=131
'^'=132
'>='=133
'>'=134
'<='=135
'<'=136
'=='=137
'='=138
'!='=139
'|'=140
'('=141
')'=142
','=143
';'=144
'->'=145
'['=146
']'=147
'?'=148
':'=149
'{'=150
'}'=151
more_math/Parser/MathExprLexer.interp
+13 -1
View File
File diff suppressed because one or more lines are too long
more_math/Parser/MathExprLexer.py
+634 -601
View File
File diff suppressed because it is too large Load Diff
more_math/Parser/MathExprLexer.tokens
+64 -58
View File
@@ -119,38 +119,42 @@ TIMESTAMP=118
NONE=119
BREAK=120
CONTINUE=121
TENSOR=122
PLUS=123
MINUS=124
MULT=125
DIV=126
MOD=127
POW=128
GE=129
GT=130
LE=131
LT=132
EQ=133
EQUEALS=134
NE=135
PIPE=136
LPAREN=137
RPAREN=138
COMMA=139
SEMICOLON=140
ARROW=141
LBRACKET=142
RBRACKET=143
QUESTION=144
COLON=145
LBRACE=146
RBRACE=147
NUMBER=148
CONSTANT=149
VARIABLE=150
SL_COMMENT=151
ML_COMMENT=152
WS=153
RIFE=122
MOTION_MASK=123
FLOW_TO_IMAGE=124
FLOW_APPLY=125
TENSOR=126
PLUS=127
MINUS=128
MULT=129
DIV=130
MOD=131
POW=132
GE=133
GT=134
LE=135
LT=136
EQ=137
EQUEALS=138
NE=139
PIPE=140
LPAREN=141
RPAREN=142
COMMA=143
SEMICOLON=144
ARROW=145
LBRACKET=146
RBRACKET=147
QUESTION=148
COLON=149
LBRACE=150
RBRACE=151
NUMBER=152
CONSTANT=153
VARIABLE=154
SL_COMMENT=155
ML_COMMENT=156
WS=157
'sin'=1
'cos'=2
'tan'=3
@@ -242,29 +246,31 @@ WS=153
'flatten'=115
'break'=120
'continue'=121
'tensor'=122
'+'=123
'-'=124
'*'=125
'/'=126
'%'=127
'^'=128
'>='=129
'>'=130
'<='=131
'<'=132
'=='=133
'='=134
'!='=135
'|'=136
'('=137
')'=138
','=139
';'=140
'->'=141
'['=142
']'=143
'?'=144
':'=145
'{'=146
'}'=147
'rife'=122
'motion_mask'=123
'tensor'=126
'+'=127
'-'=128
'*'=129
'/'=130
'%'=131
'^'=132
'>='=133
'>'=134
'<='=135
'<'=136
'=='=137
'='=138
'!='=139
'|'=140
'('=141
')'=142
','=143
';'=144
'->'=145
'['=146
']'=147
'?'=148
':'=149
'{'=150
'}'=151
more_math/Parser/MathExprListener.py
+36
View File
@@ -1088,6 +1088,24 @@ class MathExprListener(ParseTreeListener):
pass
# Enter a parse tree produced by MathExprParser#MotionMaskFunc.
def enterMotionMaskFunc(self, ctx:MathExprParser.MotionMaskFuncContext):
pass
# Exit a parse tree produced by MathExprParser#MotionMaskFunc.
def exitMotionMaskFunc(self, ctx:MathExprParser.MotionMaskFuncContext):
pass
# Enter a parse tree produced by MathExprParser#FlowToImageFunc.
def enterFlowToImageFunc(self, ctx:MathExprParser.FlowToImageFuncContext):
pass
# Exit a parse tree produced by MathExprParser#FlowToImageFunc.
def exitFlowToImageFunc(self, ctx:MathExprParser.FlowToImageFuncContext):
pass
# Enter a parse tree produced by MathExprParser#PowFunc.
def enterPowFunc(self, ctx:MathExprParser.PowFuncContext):
pass
@@ -1313,6 +1331,24 @@ class MathExprListener(ParseTreeListener):
pass
# Enter a parse tree produced by MathExprParser#FlowApplyFunc.
def enterFlowApplyFunc(self, ctx:MathExprParser.FlowApplyFuncContext):
pass
# Exit a parse tree produced by MathExprParser#FlowApplyFunc.
def exitFlowApplyFunc(self, ctx:MathExprParser.FlowApplyFuncContext):
pass
# Enter a parse tree produced by MathExprParser#RifeFunc.
def enterRifeFunc(self, ctx:MathExprParser.RifeFuncContext):
pass
# Exit a parse tree produced by MathExprParser#RifeFunc.
def exitRifeFunc(self, ctx:MathExprParser.RifeFuncContext):
pass
# Enter a parse tree produced by MathExprParser#ClampFunc.
def enterClampFunc(self, ctx:MathExprParser.ClampFuncContext):
pass
more_math/Parser/MathExprParser.py
+1064 -819
View File
File diff suppressed because it is too large Load Diff
more_math/Parser/MathExprVisitor.py
+20
View File
@@ -609,6 +609,16 @@ class MathExprVisitor(ParseTreeVisitor):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#MotionMaskFunc.
def visitMotionMaskFunc(self, ctx:MathExprParser.MotionMaskFuncContext):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#FlowToImageFunc.
def visitFlowToImageFunc(self, ctx:MathExprParser.FlowToImageFuncContext):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#PowFunc.
def visitPowFunc(self, ctx:MathExprParser.PowFuncContext):
return self.visitChildren(ctx)
@@ -734,6 +744,16 @@ class MathExprVisitor(ParseTreeVisitor):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#FlowApplyFunc.
def visitFlowApplyFunc(self, ctx:MathExprParser.FlowApplyFuncContext):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#RifeFunc.
def visitRifeFunc(self, ctx:MathExprParser.RifeFuncContext):
return self.visitChildren(ctx)
# Visit a parse tree produced by MathExprParser#ClampFunc.
def visitClampFunc(self, ctx:MathExprParser.ClampFuncContext):
return self.visitChildren(ctx)
more_math/Parser/UnifiedMathVisitor.py
+38 -1
View File
@@ -3,6 +3,7 @@ import torch
import math
import inspect
import torch.nn.functional as F
from . import optical_flow_utils as ofu
from antlr4 import TerminalNode
from .MathExprVisitor import MathExprVisitor
from ..helper_functions import generate_dim_variables
@@ -84,7 +85,7 @@ class UnifiedMathVisitor(MathExprVisitor):
return last_result.value
return last_result
def _is_tensor(self, val):
return isinstance(val, torch.Tensor)
return isinstance(val, torch.Tensor) or getattr(val, "is_nested", False)
def _is_list(self, val):
return isinstance(val, (list, tuple))
@@ -896,8 +897,15 @@ class UnifiedMathVisitor(MathExprVisitor):
def visitPermuteFunc(self, ctx):
tsr = self._promote_to_tensor((yield ctx.expr(0)))
dims = (yield ctx.expr(1))
# Ensure dims is a list of integers
if isinstance(dims, torch.Tensor):
dims = dims.flatten().long().tolist()
elif isinstance(dims, (list, tuple)):
dims = [int(0.5 + float(d)) for d in dims] # Round floats to safe ints
elif isinstance(dims, (int, float)):
dims = [int(0.5 + float(dims))]
return tsr.permute(*dims)
def visitReshapeFunc(self, ctx):
@@ -1171,6 +1179,35 @@ class UnifiedMathVisitor(MathExprVisitor):
b = self._promote_to_tensor((yield ctx.expr(1))).float()
return F.cosine_similarity(a, b, dim=-1)
def visitRifeFunc(self, ctx):
img1 = self._promote_to_tensor((yield ctx.expr(0)))
img2 = self._promote_to_tensor((yield ctx.expr(1)))
tiling_size = 0
iterations = 12
multi_scale = False
if len(ctx.expr()) >= 3:
tiling_size = float((yield ctx.expr(2)))
if len(ctx.expr()) >= 4:
iterations = int((yield ctx.expr(3)))
if len(ctx.expr()) >= 5:
multi_scale = bool((yield ctx.expr(4)))
return ofu.get_optical_flow(img1, img2, tiling_size, iterations, multi_scale)
def visitMotionMaskFunc(self, ctx):
flow = self._promote_to_tensor((yield ctx.expr()))
return ofu.calculate_occlusion_mask(flow)
def visitFlowToImageFunc(self, ctx):
flow = self._promote_to_tensor((yield ctx.expr()))
return ofu.flow_to_image(flow)
def visitFlowApplyFunc(self, ctx):
image = self._promote_to_tensor((yield ctx.expr(0)))
flow = self._promote_to_tensor((yield ctx.expr(1)))
return ofu.apply_flow(image, flow)
def visitFlipFunc(self, ctx):
val = self._promote_to_tensor((yield ctx.expr(0)))
dims = (yield ctx.expr(1))
more_math/Parser/optical_flow_utils.py
+337
View File
@@ -0,0 +1,337 @@
import torch
import torch.nn.functional as F
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights
from torchvision.utils import flow_to_image as torchvision_flow_to_image
import numpy as np
_raft_model = None
def get_raft_model(device):
global _raft_model
if _raft_model is None:
weights = Raft_Large_Weights.DEFAULT
_raft_model = raft_large(weights=weights, progress=False).to(device).eval()
return _raft_model
def warp(x, flow):
"""
Warps an image or tensor using the given flow.
x: [B, C, H, W]
flow: [B, 2, H, W] (flow in pixels)
"""
b, c, h, w = x.shape
device = flow.device # Use flow device as the master device
x = x.to(device)
# Create grid
grid_y, grid_x = torch.meshgrid(
torch.linspace(0, h - 1, h, device=device),
torch.linspace(0, w - 1, w, device=device),
indexing='ij'
)
# [H, W]
grid = torch.stack((grid_x, grid_y), dim=0).unsqueeze(0).repeat(b, 1, 1, 1) # [B, 2, H, W]
# Add flow to grid
v_grid = grid + flow # [B, 2, H, W]
# Normalize to [-1, 1] for grid_sample
v_grid_x = 2.0 * v_grid[:, 0, :, :] / max(w - 1, 1) - 1.0
v_grid_y = 2.0 * v_grid[:, 1, :, :] / max(h - 1, 1) - 1.0
v_grid = torch.stack((v_grid_x, v_grid_y), dim=3) # [B, H, W, 2]
return F.grid_sample(x, v_grid, mode='bilinear', padding_mode='reflection', align_corners=True)
def preprocess_image(img, device):
"""
img: [B, H, W, C] in range [0, 1]
Returns: [B, 3, H, W] normalized for RAFT
"""
if img.ndim == 3:
img = img.unsqueeze(0)
# ComfyUI [B, H, W, C] -> Torch [B, C, H, W]
img = img.movedim(-1, 1)
# Ensure RGB
if img.shape[1] == 1:
img = img.expand(-1, 3, -1, -1)
elif img.shape[1] > 3:
img = img[:, :3, :, :]
# Scale to [0, 255] as RAFT transforms usually expect this
img = img * 255.0
# Use official transform if possible
weights = Raft_Large_Weights.DEFAULT
transform = weights.transforms()
# The transform expects [0, 255] and returns normalized [-1, 1]
# It takes (img1, img2) but we can use it for one or just follow its logic
# Actually, let's just follow the logic: 2 * (img / 255.0) - 1.0
# Wait, if I do that, it's just 2 * img_01 - 1.0.
return (2.0 * (img / 255.0) - 1.0).to(device)
def get_optical_flow(img1, img2, tiling_size=0, iterations=12, multi_scale=False):
"""
img1, img2: [H, W, C] or [B, H, W, C]
Returns: [B, H, W, 2] flow vectors
"""
device = img1.device
model = get_raft_model(device)
img1_proc = preprocess_image(img1, device)
img2_proc = preprocess_image(img2, device)
b, _, h, w = img1_proc.shape
# Auto-tiling for large images to prevent OOM
# 2MPx is a reasonable limit for 8GB VRAM without tiling
if tiling_size <= 0 and (h * w) > 2000000:
tiling_size = 1024
# Handle fraction or float conversion
if 0 < tiling_size < 1:
tiling_size = int(min(h, w) * tiling_size)
else:
tiling_size = int(tiling_size)
# Ensure divisible by 8 for RAFT
h_new, w_new = (h // 8) * 8, (w // 8) * 8
if h != h_new or w != w_new:
img1_proc = F.interpolate(img1_proc, size=(h_new, w_new), mode='bilinear', align_corners=False)
img2_proc = F.interpolate(img2_proc, size=(h_new, w_new), mode='bilinear', align_corners=False)
# Process sequentially to save memory
flows = []
with torch.no_grad():
for i in range(b):
p1 = img1_proc[i:i+1]
p2 = img2_proc[i:i+1]
global_flow = None
if multi_scale:
# 1. Global Pass
# Downsample more aggressively to capture massive jumps (max 256)
g_size = 256
g_h = (min(h_new, g_size) // 8) * 8
g_w = (min(w_new, g_size) // 8) * 8
gp1 = F.interpolate(p1, size=(g_h, g_w), mode='bilinear', align_corners=False)
gp2 = F.interpolate(p2, size=(g_h, g_w), mode='bilinear', align_corners=False)
# Global flow estimation
global_flow = model(gp1, gp2, num_flow_updates=iterations)[-1]
# Upsample global flow to high-res
global_flow = F.interpolate(global_flow, size=(h_new, w_new), mode='bilinear', align_corners=False)
global_flow[:, 0] *= float(w_new) / g_w
global_flow[:, 1] *= float(h_new) / g_h
# print(f"DEBUG: Global flow at center: {global_flow[0, :, h_new//2, w_new//2]}")
# 2. Warp img2 using global flow
p2_warped = warp(p2, global_flow)
target_p2 = p2_warped
else:
target_p2 = p2
# 3. Local Refinement (Tiled or Full)
if tiling_size > 0:
f_residual = _tiled_flow(model, p1, target_p2, tiling_size, iterations)
else:
f_residual = model(p1, target_p2, num_flow_updates=iterations)[-1]
# 4. Combine
if global_flow is not None:
total_flow = global_flow + f_residual
else:
total_flow = f_residual
flows.append(total_flow)
if i % 4 == 0 and b > 4:
torch.cuda.empty_cache()
flow = torch.cat(flows, dim=0)
# Resize flow back if needed
if h != h_new or w != w_new:
flow = F.interpolate(flow, size=(h, w), mode='bilinear', align_corners=False)
# Rescale flow values
flow[:, 0] *= float(w) / w_new
flow[:, 1] *= float(h) / h_new
# [B, 2, H, W] -> [B, H, W, 2]
return flow.movedim(1, -1)
def calculate_occlusion_mask(flow):
"""
Calculates a disocclusion mask: pixels revealed by motion that have no
valid source data (e.g. background uncovered when an object moves away).
flow: [B, H, W, 2]
Returns: [B, H, W] mask (1.0 = disoccluded/revealed, 0.0 = valid)
"""
if flow.ndim == 3:
flow = flow.unsqueeze(0)
b, h, w, _ = flow.shape
device = flow.device
# [B, H, W, 2] -> [B, 2, H, W] for warp()
flow_nchw = flow.movedim(-1, 1)
# Warp a solid-ones image with the flow.
# Where the warped result < 1.0, those pixels sampled from outside the
# image boundary or from a region that moved away -> disoccluded.
ones = torch.ones((b, 1, h, w), device=device)
warped_ones = warp(ones, flow_nchw) # [B, 1, H, W]
# Pixels that are < threshold are disoccluded
mask = 1.0 - warped_ones.squeeze(1).clamp(0.0, 1.0) # [B, H, W]
return mask
def apply_flow(image, flow):
"""
Warps an image using optical flow vectors.
image: [B, H, W, C] (ComfyUI format)
flow: [B, H, W, 2] (from rife())
Returns: [B, H, W, C] warped image
"""
if image.ndim == 3:
image = image.unsqueeze(0)
if flow.ndim == 3:
flow = flow.unsqueeze(0)
# [B, H, W, 2] -> [B, 2, H, W]
flow_nchw = flow.movedim(-1, 1)
# Ensure image is on the same device as flow
img_nchw = image.movedim(-1, 1).to(flow.device)
warped = warp(img_nchw, flow_nchw)
# [B, C, H, W] -> [B, H, W, C]
return warped.movedim(1, -1)
def flow_to_image(flow):
"""
Converts flow vectors [B, H, W, 2] to RGB images [B, H, W, 3] for visualization.
"""
if flow.ndim == 3:
flow = flow.unsqueeze(0)
# [B, H, W, 2] -> [B, 2, H, W]
flow_torch = flow.movedim(-1, 1)
# torchvision_flow_to_image expects (N, 2, H, W)
img_uint8 = torchvision_flow_to_image(flow_torch)
# [B, 3, H, W] -> [B, H, W, 3]
img = img_uint8.movedim(1, -1).float() / 255.0
return img
def _tiled_flow(model, img1, img2, tiling_size, iterations):
"""
Robust tiling implementation with weighted blending and edge coverage.
"""
b, c, h, w = img1.shape
device = img1.device
# Enforce minimum tile size for RAFT
if tiling_size < 256:
# If user provided a very small value, maybe they meant fraction?
# But for ComfyUI, we'll just enforce a safe minimum.
tiling_size = max(256, tiling_size)
# Ensure tiling_size is divisible by 8
tiling_size = (tiling_size // 8) * 8
# Clip tiling_size to image dimensions
t_h = int(min(tiling_size, h))
t_w = int(min(tiling_size, w))
# If image is smaller than tiling_size after divisibility adjustments
if t_h < 8 or t_w < 8:
return model(img1, img2, num_flow_updates=iterations)[-1]
output_flow = torch.zeros((b, 2, h, w), device=device)
weights = torch.zeros((b, 1, h, w), device=device)
# Stride of 50% overlap is usually good
stride_h = int(t_h // 2)
stride_w = int(t_w // 2)
# Generate coordinates for tiles
y_coords = list(range(0, h - t_h, stride_h)) + [max(0, h - t_h)]
x_coords = list(range(0, w - t_w, stride_w)) + [max(0, w - t_w)]
# Remove duplicates if any (when image size matches stride)
y_coords = sorted(list(set(y_coords)))
x_coords = sorted(list(set(x_coords)))
for y in y_coords:
for x in x_coords:
patch1 = img1[:, :, y:y+t_h, x:x+t_w]
patch2 = img2[:, :, y:y+t_h, x:x+t_w]
# RAFT expects inputs divisible by 8, but we already ensured t_h, t_w are.
patch_flow = model(patch1, patch2, num_flow_updates=iterations)[-1]
# Create local mask for this tile
# Linear tapering only on overlapping edges
layer_mask = torch.ones((1, 1, t_h, t_w), device=device)
overlap_h = stride_h
overlap_w = stride_w
if y > 0: # Fade in top
layer_mask[:, :, :overlap_h, :] *= torch.linspace(0, 1, overlap_h, device=device).view(1, 1, overlap_h, 1)
if y + t_h < h: # Fade out bottom
layer_mask[:, :, -overlap_h:, :] *= torch.linspace(1, 0, overlap_h, device=device).view(1, 1, overlap_h, 1)
if x > 0: # Fade in left
layer_mask[:, :, :, :overlap_w] *= torch.linspace(0, 1, overlap_w, device=device).view(1, 1, 1, overlap_w)
if x + t_w < w: # Fade out right
layer_mask[:, :, :, -overlap_w:] *= torch.linspace(1, 0, overlap_w, device=device).view(1, 1, 1, overlap_w)
output_flow[:, :, y:y+t_h, x:x+t_w] += patch_flow * layer_mask
weights[:, :, y:y+t_h, x:x+t_w] += layer_mask
return output_flow / torch.clamp(weights, min=1e-6)
def calculate_occlusion_mask(flow):
"""
Calculates occlusion mask using flow divergence.
flow: [B, H, W, 2]
Returns: [B, H, W] mask (1.0 = occluded)
"""
b, h, w, _ = flow.shape
device = flow.device
# [B, H, W, 2] -> [B, 2, H, W]
u = flow[..., 0]
v = flow[..., 1]
# Gradient of flow
du_dx = torch.zeros_like(u)
du_dx[:, :, 1:-1] = (u[:, :, 2:] - u[:, :, :-2]) / 2.0
dv_dy = torch.zeros_like(v)
dv_dy[:, 1:-1, :] = (v[:, 2:, :] - v[:, :-2, :]) / 2.0
# Divergence
div = du_dx + dv_dy
# Occlusion heuristic: negative divergence indicates compression (potential occlusion)
# and positive divergence indicates expansion.
# We take the absolute divergence and threshold it.
mask = torch.abs(div)
# Normalize to 0-1
mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-6)
return mask