File size: 3,693 Bytes
5a6434e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Multi-Scale Flow-Warp-Mask U-Net: predicts flow at multiple resolutions."""
import torch
import torch.nn as nn
import torch.nn.functional as F
class ResConvBlock(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
self.gn1 = nn.GroupNorm(min(8, out_ch), out_ch)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
self.gn2 = nn.GroupNorm(min(8, out_ch), out_ch)
self.proj = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
def forward(self, x):
residual = self.proj(x)
x = F.silu(self.gn1(self.conv1(x)))
x = F.silu(self.gn2(self.conv2(x)))
return x + residual
class MultiScaleFlowUNet(nn.Module):
def __init__(self, in_channels=12, channels=[64, 128, 256]):
super().__init__()
# Encoder
self.encoders = nn.ModuleList()
self.pools = nn.ModuleList()
prev_ch = in_channels
for ch in channels:
self.encoders.append(ResConvBlock(prev_ch, ch))
self.pools.append(nn.MaxPool2d(2))
prev_ch = ch
# Bottleneck
self.bottleneck = ResConvBlock(channels[-1], channels[-1] * 2)
# Decoder
self.upconvs = nn.ModuleList()
self.decoders = nn.ModuleList()
dec_channels = list(reversed(channels))
prev_ch = channels[-1] * 2
for ch in dec_channels:
self.upconvs.append(nn.ConvTranspose2d(prev_ch, ch, 2, stride=2))
self.decoders.append(ResConvBlock(ch * 2, ch))
prev_ch = ch
# Multi-scale flow heads at each decoder level
# dec_channels = [256, 128, 64] (coarsest to finest)
# Level 0 (coarsest, 8x8): flow refinement
# Level 1 (16x16): flow refinement
# Level 2 (finest, 64x64): flow refinement + mask + gen_frame
self.flow_heads = nn.ModuleList()
for ch in dec_channels:
head = nn.Conv2d(ch, 2, 1)
nn.init.zeros_(head.weight)
nn.init.zeros_(head.bias)
self.flow_heads.append(head)
# Mask and generation heads only at finest level (level 2, 64x64)
self.mask_head = nn.Conv2d(dec_channels[-1], 1, 1)
nn.init.zeros_(self.mask_head.weight)
nn.init.zeros_(self.mask_head.bias)
self.gen_head = nn.Conv2d(dec_channels[-1], 3, 1)
def forward(self, x):
skips = []
for enc, pool in zip(self.encoders, self.pools):
x = enc(x)
skips.append(x)
x = pool(x)
x = self.bottleneck(x)
flows = [] # flow at each level, from coarsest to finest
for i, (upconv, dec, skip) in enumerate(zip(self.upconvs, self.decoders, reversed(skips))):
x = upconv(x)
x = torch.cat([x, skip], dim=1)
x = dec(x)
# Predict flow refinement at this level
flow_refine = self.flow_heads[i](x)
if i == 0:
# Coarsest level: just the flow refinement
flow = flow_refine
else:
# Upsample previous flow and add refinement
prev_flow_up = F.interpolate(flows[-1], scale_factor=2, mode='bilinear', align_corners=True)
# Scale flow values by 2 since coordinates double
prev_flow_up = prev_flow_up * 2
flow = prev_flow_up + flow_refine
flows.append(flow)
# Final level outputs
mask = torch.sigmoid(self.mask_head(x))
gen_frame = self.gen_head(x)
# flows[-1] is the finest (64x64) flow
return flows, mask, gen_frame
|