Add GIMM-VFI support (NeurIPS 2024) with single-pass arbitrary-timestep interpolation
Integrates GIMM-VFI alongside existing BIM/EMA/SGM models. Key feature: generates all intermediate frames in one forward pass (no recursive 2x passes needed for 4x/8x). - Vendor gimm_vfi_arch/ from kijai/ComfyUI-GIMM-VFI with device fixes - Two variants: RAFT-based (~80MB) and FlowFormer-based (~123MB) - Auto-download checkpoints from HuggingFace (Kijai/GIMM-VFI_safetensors) - Three new nodes: Load GIMM-VFI Model, GIMM-VFI Interpolate, GIMM-VFI Segment Interpolate - single_pass toggle: True=arbitrary timestep (default), False=recursive like other models - ds_factor parameter for high-res input downscaling Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
1
gimm_vfi_arch/generalizable_INR/raft/__init__.py
Normal file
1
gimm_vfi_arch/generalizable_INR/raft/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
from .raft import RAFT
|
||||
175
gimm_vfi_arch/generalizable_INR/raft/corr.py
Normal file
175
gimm_vfi_arch/generalizable_INR/raft/corr.py
Normal file
@@ -0,0 +1,175 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# --------------------------------------------------------
|
||||
# References:
|
||||
# amt: https://github.com/MCG-NKU/AMT
|
||||
# raft: https://github.com/princeton-vl/RAFT
|
||||
# --------------------------------------------------------
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from .utils.utils import bilinear_sampler, coords_grid
|
||||
|
||||
try:
|
||||
import alt_cuda_corr
|
||||
except:
|
||||
# alt_cuda_corr is not compiled
|
||||
pass
|
||||
|
||||
|
||||
class BidirCorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
self.corr_pyramid_T = []
|
||||
|
||||
corr = BidirCorrBlock.corr(fmap1, fmap2)
|
||||
batch, h1, w1, dim, h2, w2 = corr.shape
|
||||
corr_T = corr.clone().permute(0, 4, 5, 3, 1, 2)
|
||||
|
||||
corr = corr.reshape(batch * h1 * w1, dim, h2, w2)
|
||||
corr_T = corr_T.reshape(batch * h2 * w2, dim, h1, w1)
|
||||
|
||||
self.corr_pyramid.append(corr)
|
||||
self.corr_pyramid_T.append(corr_T)
|
||||
|
||||
for _ in range(self.num_levels - 1):
|
||||
corr = F.avg_pool2d(corr, 2, stride=2)
|
||||
corr_T = F.avg_pool2d(corr_T, 2, stride=2)
|
||||
self.corr_pyramid.append(corr)
|
||||
self.corr_pyramid_T.append(corr_T)
|
||||
|
||||
def __call__(self, coords0, coords1):
|
||||
r = self.radius
|
||||
coords0 = coords0.permute(0, 2, 3, 1)
|
||||
coords1 = coords1.permute(0, 2, 3, 1)
|
||||
assert (
|
||||
coords0.shape == coords1.shape
|
||||
), f"coords0 shape: [{coords0.shape}] is not equal to [{coords1.shape}]"
|
||||
batch, h1, w1, _ = coords0.shape
|
||||
|
||||
out_pyramid = []
|
||||
out_pyramid_T = []
|
||||
for i in range(self.num_levels):
|
||||
corr = self.corr_pyramid[i]
|
||||
corr_T = self.corr_pyramid_T[i]
|
||||
|
||||
dx = torch.linspace(-r, r, 2 * r + 1, device=coords0.device)
|
||||
dy = torch.linspace(-r, r, 2 * r + 1, device=coords0.device)
|
||||
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
|
||||
delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
|
||||
|
||||
centroid_lvl_0 = coords0.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
|
||||
centroid_lvl_1 = coords1.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
|
||||
coords_lvl_0 = centroid_lvl_0 + delta_lvl
|
||||
coords_lvl_1 = centroid_lvl_1 + delta_lvl
|
||||
|
||||
corr = bilinear_sampler(corr, coords_lvl_0)
|
||||
corr_T = bilinear_sampler(corr_T, coords_lvl_1)
|
||||
corr = corr.view(batch, h1, w1, -1)
|
||||
corr_T = corr_T.view(batch, h1, w1, -1)
|
||||
out_pyramid.append(corr)
|
||||
out_pyramid_T.append(corr_T)
|
||||
|
||||
out = torch.cat(out_pyramid, dim=-1)
|
||||
out_T = torch.cat(out_pyramid_T, dim=-1)
|
||||
return (
|
||||
out.permute(0, 3, 1, 2).contiguous().float(),
|
||||
out_T.permute(0, 3, 1, 2).contiguous().float(),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
batch, dim, ht, wd = fmap1.shape
|
||||
fmap1 = fmap1.view(batch, dim, ht * wd)
|
||||
fmap2 = fmap2.view(batch, dim, ht * wd)
|
||||
|
||||
corr = torch.matmul(fmap1.transpose(1, 2), fmap2)
|
||||
corr = corr.view(batch, ht, wd, 1, ht, wd)
|
||||
return corr / torch.sqrt(torch.tensor(dim).float())
|
||||
|
||||
|
||||
class AlternateCorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
|
||||
self.pyramid = [(fmap1, fmap2)]
|
||||
for i in range(self.num_levels):
|
||||
fmap1 = F.avg_pool2d(fmap1, 2, stride=2)
|
||||
fmap2 = F.avg_pool2d(fmap2, 2, stride=2)
|
||||
self.pyramid.append((fmap1, fmap2))
|
||||
|
||||
def __call__(self, coords):
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
B, H, W, _ = coords.shape
|
||||
dim = self.pyramid[0][0].shape[1]
|
||||
|
||||
corr_list = []
|
||||
for i in range(self.num_levels):
|
||||
r = self.radius
|
||||
fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous()
|
||||
fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()
|
||||
(corr,) = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r)
|
||||
corr_list.append(corr.squeeze(1))
|
||||
|
||||
corr = torch.stack(corr_list, dim=1)
|
||||
corr = corr.reshape(B, -1, H, W)
|
||||
return corr / torch.sqrt(torch.tensor(dim).float())
|
||||
|
||||
|
||||
class CorrBlock:
|
||||
def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
|
||||
self.num_levels = num_levels
|
||||
self.radius = radius
|
||||
self.corr_pyramid = []
|
||||
|
||||
# all pairs correlation
|
||||
corr = CorrBlock.corr(fmap1, fmap2)
|
||||
|
||||
batch, h1, w1, dim, h2, w2 = corr.shape
|
||||
corr = corr.reshape(batch * h1 * w1, dim, h2, w2)
|
||||
|
||||
self.corr_pyramid.append(corr)
|
||||
for i in range(self.num_levels - 1):
|
||||
corr = F.avg_pool2d(corr, 2, stride=2)
|
||||
self.corr_pyramid.append(corr)
|
||||
|
||||
def __call__(self, coords):
|
||||
r = self.radius
|
||||
coords = coords.permute(0, 2, 3, 1)
|
||||
batch, h1, w1, _ = coords.shape
|
||||
|
||||
out_pyramid = []
|
||||
for i in range(self.num_levels):
|
||||
corr = self.corr_pyramid[i]
|
||||
dx = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
|
||||
dy = torch.linspace(-r, r, 2 * r + 1, device=coords.device)
|
||||
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
|
||||
|
||||
centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
|
||||
delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
|
||||
coords_lvl = centroid_lvl + delta_lvl
|
||||
|
||||
corr = bilinear_sampler(corr, coords_lvl)
|
||||
corr = corr.view(batch, h1, w1, -1)
|
||||
out_pyramid.append(corr)
|
||||
|
||||
out = torch.cat(out_pyramid, dim=-1)
|
||||
return out.permute(0, 3, 1, 2).contiguous().float()
|
||||
|
||||
@staticmethod
|
||||
def corr(fmap1, fmap2):
|
||||
batch, dim, ht, wd = fmap1.shape
|
||||
fmap1 = fmap1.view(batch, dim, ht * wd)
|
||||
fmap2 = fmap2.view(batch, dim, ht * wd)
|
||||
|
||||
corr = torch.matmul(fmap1.transpose(1, 2), fmap2)
|
||||
corr = corr.view(batch, ht, wd, 1, ht, wd)
|
||||
return corr / torch.sqrt(torch.tensor(dim).float())
|
||||
293
gimm_vfi_arch/generalizable_INR/raft/extractor.py
Normal file
293
gimm_vfi_arch/generalizable_INR/raft/extractor.py
Normal file
@@ -0,0 +1,293 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn="group", stride=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_planes, planes, kernel_size=3, padding=1, stride=stride
|
||||
)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == "group":
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == "batch":
|
||||
self.norm1 = nn.BatchNorm2d(planes)
|
||||
self.norm2 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == "instance":
|
||||
self.norm1 = nn.InstanceNorm2d(planes)
|
||||
self.norm2 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == "none":
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm3 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x + y)
|
||||
|
||||
|
||||
class BottleneckBlock(nn.Module):
|
||||
def __init__(self, in_planes, planes, norm_fn="group", stride=1):
|
||||
super(BottleneckBlock, self).__init__()
|
||||
|
||||
self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0)
|
||||
self.conv2 = nn.Conv2d(
|
||||
planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride
|
||||
)
|
||||
self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
num_groups = planes // 8
|
||||
|
||||
if norm_fn == "group":
|
||||
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4)
|
||||
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4)
|
||||
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
|
||||
|
||||
elif norm_fn == "batch":
|
||||
self.norm1 = nn.BatchNorm2d(planes // 4)
|
||||
self.norm2 = nn.BatchNorm2d(planes // 4)
|
||||
self.norm3 = nn.BatchNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.BatchNorm2d(planes)
|
||||
|
||||
elif norm_fn == "instance":
|
||||
self.norm1 = nn.InstanceNorm2d(planes // 4)
|
||||
self.norm2 = nn.InstanceNorm2d(planes // 4)
|
||||
self.norm3 = nn.InstanceNorm2d(planes)
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.InstanceNorm2d(planes)
|
||||
|
||||
elif norm_fn == "none":
|
||||
self.norm1 = nn.Sequential()
|
||||
self.norm2 = nn.Sequential()
|
||||
self.norm3 = nn.Sequential()
|
||||
if not stride == 1:
|
||||
self.norm4 = nn.Sequential()
|
||||
|
||||
if stride == 1:
|
||||
self.downsample = None
|
||||
|
||||
else:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
y = x
|
||||
y = self.relu(self.norm1(self.conv1(y)))
|
||||
y = self.relu(self.norm2(self.conv2(y)))
|
||||
y = self.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
if self.downsample is not None:
|
||||
x = self.downsample(x)
|
||||
|
||||
return self.relu(x + y)
|
||||
|
||||
|
||||
class BasicEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0, only_feat=False):
|
||||
super(BasicEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
self.only_feat = only_feat
|
||||
|
||||
if self.norm_fn == "group":
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
|
||||
|
||||
elif self.norm_fn == "batch":
|
||||
self.norm1 = nn.BatchNorm2d(64)
|
||||
|
||||
elif self.norm_fn == "instance":
|
||||
self.norm1 = nn.InstanceNorm2d(64)
|
||||
|
||||
elif self.norm_fn == "none":
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 64
|
||||
self.layer1 = self._make_layer(64, stride=1)
|
||||
self.layer2 = self._make_layer(96, stride=2)
|
||||
self.layer3 = self._make_layer(128, stride=2)
|
||||
|
||||
if not self.only_feat:
|
||||
# output convolution
|
||||
self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
|
||||
|
||||
self.dropout = None
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, return_feature=False, mif=False):
|
||||
features = []
|
||||
|
||||
# if input is list, combine batch dimension
|
||||
is_list = isinstance(x, tuple) or isinstance(x, list)
|
||||
if is_list:
|
||||
batch_dim = x[0].shape[0]
|
||||
x = torch.cat(x, dim=0)
|
||||
|
||||
x_2 = F.interpolate(x, scale_factor=1 / 2, mode="bilinear", align_corners=False)
|
||||
x_4 = F.interpolate(x, scale_factor=1 / 4, mode="bilinear", align_corners=False)
|
||||
|
||||
def f1(feat):
|
||||
feat = self.conv1(feat)
|
||||
feat = self.norm1(feat)
|
||||
feat = self.relu1(feat)
|
||||
feat = self.layer1(feat)
|
||||
return feat
|
||||
|
||||
x = f1(x)
|
||||
features.append(x)
|
||||
x = self.layer2(x)
|
||||
if mif:
|
||||
x_2_2 = f1(x_2)
|
||||
features.append(torch.cat([x, x_2_2], dim=1))
|
||||
else:
|
||||
features.append(x)
|
||||
x = self.layer3(x)
|
||||
if mif:
|
||||
x_2_4 = self.layer2(x_2_2)
|
||||
x_4_4 = f1(x_4)
|
||||
features.append(torch.cat([x, x_2_4, x_4_4], dim=1))
|
||||
else:
|
||||
features.append(x)
|
||||
|
||||
if not self.only_feat:
|
||||
x = self.conv2(x)
|
||||
|
||||
if self.training and self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
|
||||
if is_list:
|
||||
x = torch.split(x, [batch_dim, batch_dim], dim=0)
|
||||
features = [torch.split(f, [batch_dim, batch_dim], dim=0) for f in features]
|
||||
if return_feature:
|
||||
return x, features
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class SmallEncoder(nn.Module):
|
||||
def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0):
|
||||
super(SmallEncoder, self).__init__()
|
||||
self.norm_fn = norm_fn
|
||||
|
||||
if self.norm_fn == "group":
|
||||
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)
|
||||
|
||||
elif self.norm_fn == "batch":
|
||||
self.norm1 = nn.BatchNorm2d(32)
|
||||
|
||||
elif self.norm_fn == "instance":
|
||||
self.norm1 = nn.InstanceNorm2d(32)
|
||||
|
||||
elif self.norm_fn == "none":
|
||||
self.norm1 = nn.Sequential()
|
||||
|
||||
self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
|
||||
self.in_planes = 32
|
||||
self.layer1 = self._make_layer(32, stride=1)
|
||||
self.layer2 = self._make_layer(64, stride=2)
|
||||
self.layer3 = self._make_layer(96, stride=2)
|
||||
|
||||
self.dropout = None
|
||||
if dropout > 0:
|
||||
self.dropout = nn.Dropout2d(p=dropout)
|
||||
|
||||
self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
|
||||
if m.weight is not None:
|
||||
nn.init.constant_(m.weight, 1)
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_layer(self, dim, stride=1):
|
||||
layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)
|
||||
layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)
|
||||
layers = (layer1, layer2)
|
||||
|
||||
self.in_planes = dim
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
# if input is list, combine batch dimension
|
||||
is_list = isinstance(x, tuple) or isinstance(x, list)
|
||||
if is_list:
|
||||
batch_dim = x[0].shape[0]
|
||||
x = torch.cat(x, dim=0)
|
||||
|
||||
x = self.conv1(x)
|
||||
x = self.norm1(x)
|
||||
x = self.relu1(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.conv2(x)
|
||||
|
||||
if self.training and self.dropout is not None:
|
||||
x = self.dropout(x)
|
||||
|
||||
if is_list:
|
||||
x = torch.split(x, [batch_dim, batch_dim], dim=0)
|
||||
|
||||
return x
|
||||
238
gimm_vfi_arch/generalizable_INR/raft/other_raft.py
Normal file
238
gimm_vfi_arch/generalizable_INR/raft/other_raft.py
Normal file
@@ -0,0 +1,238 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .update import BasicUpdateBlock, SmallUpdateBlock
|
||||
from .extractor import BasicEncoder, SmallEncoder
|
||||
from .corr import BidirCorrBlock, AlternateCorrBlock
|
||||
from .utils.utils import bilinear_sampler, coords_grid, upflow8
|
||||
|
||||
try:
|
||||
autocast = torch.cuda.amp.autocast
|
||||
except:
|
||||
# dummy autocast for PyTorch < 1.6
|
||||
class autocast:
|
||||
def __init__(self, enabled):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
# BiRAFT
|
||||
class RAFT(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(RAFT, self).__init__()
|
||||
self.args = args
|
||||
|
||||
if args.small:
|
||||
self.hidden_dim = hdim = 96
|
||||
self.context_dim = cdim = 64
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 3
|
||||
self.corr_levels = 4
|
||||
self.corr_radius = 3
|
||||
|
||||
else:
|
||||
self.hidden_dim = hdim = 128
|
||||
self.context_dim = cdim = 128
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 4
|
||||
self.corr_levels = 4
|
||||
self.corr_radius = 4
|
||||
|
||||
if "dropout" not in args._get_kwargs():
|
||||
self.args.dropout = 0
|
||||
|
||||
if "alternate_corr" not in args._get_kwargs():
|
||||
self.args.alternate_corr = False
|
||||
|
||||
# feature network, context network, and update block
|
||||
if args.small:
|
||||
self.fnet = SmallEncoder(
|
||||
output_dim=128, norm_fn="instance", dropout=args.dropout
|
||||
)
|
||||
self.cnet = SmallEncoder(
|
||||
output_dim=hdim + cdim, norm_fn="none", dropout=args.dropout
|
||||
)
|
||||
self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
else:
|
||||
self.fnet = BasicEncoder(
|
||||
output_dim=256, norm_fn="instance", dropout=args.dropout
|
||||
)
|
||||
self.cnet = BasicEncoder(
|
||||
output_dim=hdim + cdim, norm_fn="batch", dropout=args.dropout
|
||||
)
|
||||
self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
def freeze_bn(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.BatchNorm2d):
|
||||
m.eval()
|
||||
|
||||
def build_coord(self, img):
|
||||
N, C, H, W = img.shape
|
||||
coords = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
return coords
|
||||
|
||||
def initialize_flow(self, img, img2):
|
||||
"""Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
|
||||
assert img.shape == img2.shape
|
||||
N, C, H, W = img.shape
|
||||
coords01 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
coords02 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
coords1 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
coords2 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
|
||||
# optical flow computed as difference: flow = coords1 - coords0
|
||||
return coords01, coords02, coords1, coords2
|
||||
|
||||
def upsample_flow(self, flow, mask):
|
||||
"""Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination"""
|
||||
N, _, H, W = flow.shape
|
||||
mask = mask.view(N, 1, 9, 8, 8, H, W)
|
||||
mask = torch.softmax(mask, dim=2)
|
||||
|
||||
up_flow = F.unfold(8 * flow, [3, 3], padding=1)
|
||||
up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
|
||||
|
||||
up_flow = torch.sum(mask * up_flow, dim=2)
|
||||
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
|
||||
return up_flow.reshape(N, 2, 8 * H, 8 * W)
|
||||
|
||||
def get_corr_fn(self, image1, image2, projector=None):
|
||||
# run the feature network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
fmaps, feats = self.fnet([image1, image2], return_feature=True)
|
||||
fmap1, fmap2 = fmaps
|
||||
fmap1 = fmap1.float()
|
||||
fmap2 = fmap2.float()
|
||||
corr_fn1 = None
|
||||
if self.args.alternate_corr:
|
||||
corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
if projector is not None:
|
||||
corr_fn1 = AlternateCorrBlock(
|
||||
projector(feats[-1][0]),
|
||||
projector(feats[-1][1]),
|
||||
radius=self.args.corr_radius,
|
||||
)
|
||||
else:
|
||||
corr_fn = BidirCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
if projector is not None:
|
||||
corr_fn1 = BidirCorrBlock(
|
||||
projector(feats[-1][0]),
|
||||
projector(feats[-1][1]),
|
||||
radius=self.args.corr_radius,
|
||||
)
|
||||
if corr_fn1 is None:
|
||||
return corr_fn, corr_fn
|
||||
else:
|
||||
return corr_fn, corr_fn1
|
||||
|
||||
def get_corr_fn_from_feat(self, fmap1, fmap2):
|
||||
fmap1 = fmap1.float()
|
||||
fmap2 = fmap2.float()
|
||||
if self.args.alternate_corr:
|
||||
corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
else:
|
||||
corr_fn = BidirCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
|
||||
return corr_fn
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image1,
|
||||
image2,
|
||||
iters=12,
|
||||
flow_init=None,
|
||||
upsample=True,
|
||||
test_mode=False,
|
||||
corr_fn=None,
|
||||
mif=False,
|
||||
):
|
||||
"""Estimate optical flow between pair of frames"""
|
||||
assert flow_init is None
|
||||
|
||||
image1 = 2 * (image1 / 255.0) - 1.0
|
||||
image2 = 2 * (image2 / 255.0) - 1.0
|
||||
|
||||
image1 = image1.contiguous()
|
||||
image2 = image2.contiguous()
|
||||
|
||||
hdim = self.hidden_dim
|
||||
cdim = self.context_dim
|
||||
|
||||
if corr_fn is None:
|
||||
corr_fn, _ = self.get_corr_fn(image1, image2)
|
||||
|
||||
# # run the feature network
|
||||
# with autocast(enabled=self.args.mixed_precision):
|
||||
# fmap1, fmap2 = self.fnet([image1, image2])
|
||||
|
||||
# fmap1 = fmap1.float()
|
||||
# fmap2 = fmap2.float()
|
||||
# if self.args.alternate_corr:
|
||||
# corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
# else:
|
||||
# corr_fn = BidirCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
|
||||
# run the context network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
# for image1
|
||||
cnet1, features1 = self.cnet(image1, return_feature=True, mif=mif)
|
||||
net1, inp1 = torch.split(cnet1, [hdim, cdim], dim=1)
|
||||
net1 = torch.tanh(net1)
|
||||
inp1 = torch.relu(inp1)
|
||||
# for image2
|
||||
cnet2, features2 = self.cnet(image2, return_feature=True, mif=mif)
|
||||
net2, inp2 = torch.split(cnet2, [hdim, cdim], dim=1)
|
||||
net2 = torch.tanh(net2)
|
||||
inp2 = torch.relu(inp2)
|
||||
|
||||
coords01, coords02, coords1, coords2 = self.initialize_flow(image1, image2)
|
||||
|
||||
# if flow_init is not None:
|
||||
# coords1 = coords1 + flow_init
|
||||
|
||||
# flow_predictions1 = []
|
||||
# flow_predictions2 = []
|
||||
for itr in range(iters):
|
||||
coords1 = coords1.detach()
|
||||
coords2 = coords2.detach()
|
||||
corr1, corr2 = corr_fn(coords1, coords2) # index correlation volume
|
||||
|
||||
flow1 = coords1 - coords01
|
||||
flow2 = coords2 - coords02
|
||||
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
net1, up_mask1, delta_flow1 = self.update_block(
|
||||
net1, inp1, corr1, flow1
|
||||
)
|
||||
net2, up_mask2, delta_flow2 = self.update_block(
|
||||
net2, inp2, corr2, flow2
|
||||
)
|
||||
|
||||
# F(t+1) = F(t) + \Delta(t)
|
||||
coords1 = coords1 + delta_flow1
|
||||
coords2 = coords2 + delta_flow2
|
||||
flow_low1 = coords1 - coords01
|
||||
flow_low2 = coords2 - coords02
|
||||
# upsample predictions
|
||||
if up_mask1 is None:
|
||||
flow_up1 = upflow8(coords1 - coords01)
|
||||
flow_up2 = upflow8(coords2 - coords02)
|
||||
else:
|
||||
flow_up1 = self.upsample_flow(coords1 - coords01, up_mask1)
|
||||
flow_up2 = self.upsample_flow(coords2 - coords02, up_mask2)
|
||||
|
||||
# flow_predictions.append(flow_up)
|
||||
return flow_up1, flow_up2, flow_low1, flow_low2, features1, features2
|
||||
# if test_mode:
|
||||
# return coords1 - coords0, flow_up
|
||||
|
||||
# return flow_predictions
|
||||
169
gimm_vfi_arch/generalizable_INR/raft/raft.py
Normal file
169
gimm_vfi_arch/generalizable_INR/raft/raft.py
Normal file
@@ -0,0 +1,169 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .update import BasicUpdateBlock, SmallUpdateBlock
|
||||
from .extractor import BasicEncoder, SmallEncoder
|
||||
from .corr import CorrBlock, AlternateCorrBlock
|
||||
from .utils.utils import bilinear_sampler, coords_grid, upflow8
|
||||
|
||||
try:
|
||||
autocast = torch.cuda.amp.autocast
|
||||
except:
|
||||
# dummy autocast for PyTorch < 1.6
|
||||
class autocast:
|
||||
def __init__(self, enabled):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
class RAFT(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(RAFT, self).__init__()
|
||||
self.args = args
|
||||
|
||||
if args.small:
|
||||
self.hidden_dim = hdim = 96
|
||||
self.context_dim = cdim = 64
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 3
|
||||
self.corr_levels = 4
|
||||
self.corr_radius = 3
|
||||
|
||||
else:
|
||||
self.hidden_dim = hdim = 128
|
||||
self.context_dim = cdim = 128
|
||||
args.corr_levels = 4
|
||||
args.corr_radius = 4
|
||||
self.corr_levels = 4
|
||||
self.corr_radius = 4
|
||||
|
||||
if "dropout" not in args._get_kwargs():
|
||||
self.args.dropout = 0
|
||||
|
||||
if "alternate_corr" not in args._get_kwargs():
|
||||
self.args.alternate_corr = False
|
||||
|
||||
# feature network, context network, and update block
|
||||
if args.small:
|
||||
self.fnet = SmallEncoder(
|
||||
output_dim=128, norm_fn="instance", dropout=args.dropout
|
||||
)
|
||||
self.cnet = SmallEncoder(
|
||||
output_dim=hdim + cdim, norm_fn="none", dropout=args.dropout
|
||||
)
|
||||
self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
else:
|
||||
self.fnet = BasicEncoder(
|
||||
output_dim=256, norm_fn="instance", dropout=args.dropout
|
||||
)
|
||||
self.cnet = BasicEncoder(
|
||||
output_dim=hdim + cdim, norm_fn="batch", dropout=args.dropout
|
||||
)
|
||||
self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
|
||||
|
||||
def freeze_bn(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.BatchNorm2d):
|
||||
m.eval()
|
||||
|
||||
def initialize_flow(self, img):
|
||||
"""Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
|
||||
N, C, H, W = img.shape
|
||||
coords0 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
coords1 = coords_grid(N, H // 8, W // 8, device=img.device)
|
||||
|
||||
# optical flow computed as difference: flow = coords1 - coords0
|
||||
return coords0, coords1
|
||||
|
||||
def upsample_flow(self, flow, mask):
|
||||
"""Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination"""
|
||||
N, _, H, W = flow.shape
|
||||
mask = mask.view(N, 1, 9, 8, 8, H, W)
|
||||
mask = torch.softmax(mask, dim=2)
|
||||
|
||||
up_flow = F.unfold(8 * flow, [3, 3], padding=1)
|
||||
up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
|
||||
|
||||
up_flow = torch.sum(mask * up_flow, dim=2)
|
||||
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
|
||||
return up_flow.reshape(N, 2, 8 * H, 8 * W)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image1,
|
||||
image2,
|
||||
iters=12,
|
||||
flow_init=None,
|
||||
upsample=True,
|
||||
test_mode=False,
|
||||
return_feat=True,
|
||||
):
|
||||
"""Estimate optical flow between pair of frames"""
|
||||
|
||||
image1 = 2 * (image1 / 255.0) - 1.0
|
||||
image2 = 2 * (image2 / 255.0) - 1.0
|
||||
|
||||
image1 = image1.contiguous()
|
||||
image2 = image2.contiguous()
|
||||
|
||||
hdim = self.hidden_dim
|
||||
cdim = self.context_dim
|
||||
|
||||
# run the feature network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
fmap1, fmap2 = self.fnet([image1, image2])
|
||||
|
||||
fmap1 = fmap1.float()
|
||||
fmap2 = fmap2.float()
|
||||
if self.args.alternate_corr:
|
||||
corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
else:
|
||||
corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
|
||||
|
||||
# run the context network
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
cnet, feats = self.cnet(image1, return_feature=True)
|
||||
net, inp = torch.split(cnet, [hdim, cdim], dim=1)
|
||||
net = torch.tanh(net)
|
||||
inp = torch.relu(inp)
|
||||
|
||||
coords0, coords1 = self.initialize_flow(image1)
|
||||
|
||||
if flow_init is not None:
|
||||
coords1 = coords1 + flow_init
|
||||
|
||||
flow_predictions = []
|
||||
for itr in range(iters):
|
||||
coords1 = coords1.detach()
|
||||
corr = corr_fn(coords1) # index correlation volume
|
||||
|
||||
flow = coords1 - coords0
|
||||
with autocast(enabled=self.args.mixed_precision):
|
||||
net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
|
||||
|
||||
# F(t+1) = F(t) + \Delta(t)
|
||||
coords1 = coords1 + delta_flow
|
||||
|
||||
# upsample predictions
|
||||
if up_mask is None:
|
||||
flow_up = upflow8(coords1 - coords0)
|
||||
else:
|
||||
flow_up = self.upsample_flow(coords1 - coords0, up_mask)
|
||||
|
||||
flow_predictions.append(flow_up)
|
||||
|
||||
if test_mode:
|
||||
return coords1 - coords0, flow_up
|
||||
|
||||
if return_feat:
|
||||
return flow_up, feats[1:], fmap1
|
||||
|
||||
return flow_predictions
|
||||
154
gimm_vfi_arch/generalizable_INR/raft/update.py
Normal file
154
gimm_vfi_arch/generalizable_INR/raft/update.py
Normal file
@@ -0,0 +1,154 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class FlowHead(nn.Module):
|
||||
def __init__(self, input_dim=128, hidden_dim=256):
|
||||
super(FlowHead, self).__init__()
|
||||
self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
|
||||
self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv2(self.relu(self.conv1(x)))
|
||||
|
||||
|
||||
class ConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim=128, input_dim=192 + 128):
|
||||
super(ConvGRU, self).__init__()
|
||||
self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
||||
self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
||||
self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
||||
|
||||
def forward(self, h, x):
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
|
||||
z = torch.sigmoid(self.convz(hx))
|
||||
r = torch.sigmoid(self.convr(hx))
|
||||
q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1)))
|
||||
|
||||
h = (1 - z) * h + z * q
|
||||
return h
|
||||
|
||||
|
||||
class SepConvGRU(nn.Module):
|
||||
def __init__(self, hidden_dim=128, input_dim=192 + 128):
|
||||
super(SepConvGRU, self).__init__()
|
||||
self.convz1 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
|
||||
)
|
||||
self.convr1 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
|
||||
)
|
||||
self.convq1 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
|
||||
)
|
||||
|
||||
self.convz2 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
|
||||
)
|
||||
self.convr2 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
|
||||
)
|
||||
self.convq2 = nn.Conv2d(
|
||||
hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
|
||||
)
|
||||
|
||||
def forward(self, h, x):
|
||||
# horizontal
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz1(hx))
|
||||
r = torch.sigmoid(self.convr1(hx))
|
||||
q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1)))
|
||||
h = (1 - z) * h + z * q
|
||||
|
||||
# vertical
|
||||
hx = torch.cat([h, x], dim=1)
|
||||
z = torch.sigmoid(self.convz2(hx))
|
||||
r = torch.sigmoid(self.convr2(hx))
|
||||
q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1)))
|
||||
h = (1 - z) * h + z * q
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class SmallMotionEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(SmallMotionEncoder, self).__init__()
|
||||
cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2
|
||||
self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0)
|
||||
self.convf1 = nn.Conv2d(2, 64, 7, padding=3)
|
||||
self.convf2 = nn.Conv2d(64, 32, 3, padding=1)
|
||||
self.conv = nn.Conv2d(128, 80, 3, padding=1)
|
||||
|
||||
def forward(self, flow, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
flo = F.relu(self.convf1(flow))
|
||||
flo = F.relu(self.convf2(flo))
|
||||
cor_flo = torch.cat([cor, flo], dim=1)
|
||||
out = F.relu(self.conv(cor_flo))
|
||||
return torch.cat([out, flow], dim=1)
|
||||
|
||||
|
||||
class BasicMotionEncoder(nn.Module):
|
||||
def __init__(self, args):
|
||||
super(BasicMotionEncoder, self).__init__()
|
||||
cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2
|
||||
self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
|
||||
self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
|
||||
self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
|
||||
self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
|
||||
self.conv = nn.Conv2d(64 + 192, 128 - 2, 3, padding=1)
|
||||
|
||||
def forward(self, flow, corr):
|
||||
cor = F.relu(self.convc1(corr))
|
||||
cor = F.relu(self.convc2(cor))
|
||||
flo = F.relu(self.convf1(flow))
|
||||
flo = F.relu(self.convf2(flo))
|
||||
|
||||
cor_flo = torch.cat([cor, flo], dim=1)
|
||||
out = F.relu(self.conv(cor_flo))
|
||||
return torch.cat([out, flow], dim=1)
|
||||
|
||||
|
||||
class SmallUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dim=96):
|
||||
super(SmallUpdateBlock, self).__init__()
|
||||
self.encoder = SmallMotionEncoder(args)
|
||||
self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82 + 64)
|
||||
self.flow_head = FlowHead(hidden_dim, hidden_dim=128)
|
||||
|
||||
def forward(self, net, inp, corr, flow):
|
||||
motion_features = self.encoder(flow, corr)
|
||||
inp = torch.cat([inp, motion_features], dim=1)
|
||||
net = self.gru(net, inp)
|
||||
delta_flow = self.flow_head(net)
|
||||
|
||||
return net, None, delta_flow
|
||||
|
||||
|
||||
class BasicUpdateBlock(nn.Module):
|
||||
def __init__(self, args, hidden_dim=128, input_dim=128):
|
||||
super(BasicUpdateBlock, self).__init__()
|
||||
self.args = args
|
||||
self.encoder = BasicMotionEncoder(args)
|
||||
self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128 + hidden_dim)
|
||||
self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
|
||||
|
||||
self.mask = nn.Sequential(
|
||||
nn.Conv2d(128, 256, 3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(256, 64 * 9, 1, padding=0),
|
||||
)
|
||||
|
||||
def forward(self, net, inp, corr, flow, upsample=True):
|
||||
motion_features = self.encoder(flow.to(inp), corr.to(inp))
|
||||
inp = torch.cat([inp, motion_features], dim=1)
|
||||
|
||||
net = self.gru(net, inp)
|
||||
delta_flow = self.flow_head(net)
|
||||
|
||||
# scale mask to balence gradients
|
||||
mask = 0.25 * self.mask(net)
|
||||
return net, mask, delta_flow
|
||||
93
gimm_vfi_arch/generalizable_INR/raft/utils/utils.py
Normal file
93
gimm_vfi_arch/generalizable_INR/raft/utils/utils.py
Normal file
@@ -0,0 +1,93 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from scipy import interpolate
|
||||
|
||||
|
||||
class InputPadder:
|
||||
"""Pads images such that dimensions are divisible by 8"""
|
||||
|
||||
def __init__(self, dims, mode="sintel"):
|
||||
self.ht, self.wd = dims[-2:]
|
||||
pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8
|
||||
pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8
|
||||
if mode == "sintel":
|
||||
self._pad = [
|
||||
pad_wd // 2,
|
||||
pad_wd - pad_wd // 2,
|
||||
pad_ht // 2,
|
||||
pad_ht - pad_ht // 2,
|
||||
]
|
||||
else:
|
||||
self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht]
|
||||
|
||||
def pad(self, *inputs):
|
||||
return [F.pad(x, self._pad, mode="replicate") for x in inputs]
|
||||
|
||||
def unpad(self, x):
|
||||
ht, wd = x.shape[-2:]
|
||||
c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]]
|
||||
return x[..., c[0] : c[1], c[2] : c[3]]
|
||||
|
||||
|
||||
def forward_interpolate(flow):
|
||||
flow = flow.detach().cpu().numpy()
|
||||
dx, dy = flow[0], flow[1]
|
||||
|
||||
ht, wd = dx.shape
|
||||
x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
|
||||
|
||||
x1 = x0 + dx
|
||||
y1 = y0 + dy
|
||||
|
||||
x1 = x1.reshape(-1)
|
||||
y1 = y1.reshape(-1)
|
||||
dx = dx.reshape(-1)
|
||||
dy = dy.reshape(-1)
|
||||
|
||||
valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
|
||||
x1 = x1[valid]
|
||||
y1 = y1[valid]
|
||||
dx = dx[valid]
|
||||
dy = dy[valid]
|
||||
|
||||
flow_x = interpolate.griddata(
|
||||
(x1, y1), dx, (x0, y0), method="nearest", fill_value=0
|
||||
)
|
||||
|
||||
flow_y = interpolate.griddata(
|
||||
(x1, y1), dy, (x0, y0), method="nearest", fill_value=0
|
||||
)
|
||||
|
||||
flow = np.stack([flow_x, flow_y], axis=0)
|
||||
return torch.from_numpy(flow).float()
|
||||
|
||||
|
||||
def bilinear_sampler(img, coords, mode="bilinear", mask=False):
|
||||
"""Wrapper for grid_sample, uses pixel coordinates"""
|
||||
H, W = img.shape[-2:]
|
||||
xgrid, ygrid = coords.split([1, 1], dim=-1)
|
||||
xgrid = 2 * xgrid / (W - 1) - 1
|
||||
ygrid = 2 * ygrid / (H - 1) - 1
|
||||
|
||||
grid = torch.cat([xgrid, ygrid], dim=-1)
|
||||
img = F.grid_sample(img, grid, align_corners=True)
|
||||
|
||||
if mask:
|
||||
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
|
||||
return img, mask.float()
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def coords_grid(batch, ht, wd, device):
|
||||
coords = torch.meshgrid(
|
||||
torch.arange(ht, device=device), torch.arange(wd, device=device)
|
||||
)
|
||||
coords = torch.stack(coords[::-1], dim=0).float()
|
||||
return coords[None].repeat(batch, 1, 1, 1)
|
||||
|
||||
|
||||
def upflow8(flow, mode="bilinear"):
|
||||
new_size = (8 * flow.shape[2], 8 * flow.shape[3])
|
||||
return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
|
||||
Reference in New Issue
Block a user