""" https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_scripts/interpolate.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/models/ssldtm.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/util_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/twodee_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/pytorch_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/distance_transform_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/sketchers_v1.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/interpolator_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/gridnet_v1.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/flow_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_util/softsplat_v0.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/rfr_new.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/extractor.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/update.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/corr.py https://github.com/ShuhongChen/eisai-anime-interpolator/blob/master/_train/frame_interpolation/helpers/raft_v1/utils.py """ import copy import cv2 import torch.nn.functional as F import torchvision.transforms.functional as F import gc from PIL import Image, ImageFile, ImageFont, ImageDraw import inspect from scipy import interpolate import kornia import math from argparse import Namespace import torch.nn as nn import numpy as np import os from functools import partial import pathlib import PIL import re import requests from scipy.spatial.transform import Rotation import scipy import shutil import torchvision.transforms as T import time import torch import torchvision as tv import zlib import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from tqdm.auto import tqdm as std_tqdm from tqdm.auto import trange as std_trange from vfi_models.ops import FunctionSoftsplat, batch_edt from comfy.model_management import get_torch_device device = get_torch_device() autocast = torch.autocast tqdm = partial(std_tqdm, dynamic_ncols=True) trange = partial(std_trange, dynamic_ncols=True) ImageFile.LOAD_TRUNCATED_IMAGES = True def pixel_ij(x, rounding=True): if isinstance(x, np.ndarray): x = x.tolist() return tuple( pixel_rounder(i, rounding) for i in (x if isinstance(x, tuple) or isinstance(x, list) else (x, x)) ) def rescale_dry(x, factor): h, w = x[-2:] if isinstance(x, tuple) or isinstance(x, list) else I(x).size return (h * factor, w * factor) def pixel_rounder(n, mode): if mode == True or mode == "round": return round(n) elif mode == "ceil": return math.ceil(n) elif mode == "floor": return math.floor(n) else: return n def diam(x): if isinstance(x, tuple) or isinstance(x, list): h, w = x[-2:] elif isinstance(x, I): h, w = x.size else: h, w = x.shape[-2:] return np.sqrt(h**2 + w**2) def pixel_logit(x, pixel_margin=1): x = (x * (255 - 2 * pixel_margin) + pixel_margin) / 255 return torch.log(x / (1 - x)) class InputPadder: """Pads images such that dimensions are divisible by 8""" def __init__(self, dims): 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 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="cubic", fill_value=0) flow_y = interpolate.griddata((x1, y1), dy, (x0, y0), method="cubic", 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) # print(img.size()) 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): coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) 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) 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) dy = torch.linspace(-r, r, 2 * r + 1) delta = torch.stack(torch.meshgrid(dy, dx), dim=-1).to(coords.device) 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()) 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, corr) 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 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): super(BasicEncoder, self).__init__() self.norm_fn = norm_fn 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) # 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): # 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 class BasicEncoder1(nn.Module): def __init__(self, output_dim=128, norm_fn="batch", dropout=0.0): super(BasicEncoder1, self).__init__() self.norm_fn = norm_fn 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(2, 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) # 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): # 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 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 ################################################## # RFR is implemented based on RAFT optical flow # ################################################## def backwarp(img, flow): _, _, H, W = img.size() u = flow[:, 0, :, :] v = flow[:, 1, :, :] gridX, gridY = np.meshgrid(np.arange(W), np.arange(H)) gridX = torch.tensor( gridX, requires_grad=False, ).cuda() gridY = torch.tensor( gridY, requires_grad=False, ).cuda() x = gridX.unsqueeze(0).expand_as(u).float() + u y = gridY.unsqueeze(0).expand_as(v).float() + v # range -1 to 1 x = 2 * (x / (W - 1) - 0.5) y = 2 * (y / (H - 1) - 0.5) # stacking X and Y grid = torch.stack((x, y), dim=3) # Sample pixels using bilinear interpolation. imgOut = torch.nn.functional.grid_sample(img, grid, align_corners=True) return imgOut class ErrorAttention(nn.Module): """A three-layer network for predicting mask""" def __init__(self, input, output): super(ErrorAttention, self).__init__() self.conv1 = nn.Conv2d(input, 32, 5, padding=2) self.conv2 = nn.Conv2d(32, 32, 3, padding=1) self.conv3 = nn.Conv2d(38, output, 3, padding=1) self.prelu1 = nn.PReLU() self.prelu2 = nn.PReLU() def forward(self, x1): x = self.prelu1(self.conv1(x1)) x = self.prelu2(torch.cat([self.conv2(x), x1], dim=1)) x = self.conv3(x) return x class RFR(nn.Module): def __init__(self, args): super(RFR, self).__init__() self.attention2 = ErrorAttention(6, 1) self.hidden_dim = hdim = 128 self.context_dim = cdim = 128 args.corr_levels = 4 args.corr_radius = 4 args.dropout = 0 self.args = args # feature network, context network, and update block self.fnet = BasicEncoder(output_dim=256, norm_fn="none", dropout=args.dropout) # self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='none', 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).to(img.device) coords1 = coords_grid(N, H // 8, W // 8).to(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 ): H, W = image1.size()[2:4] H8 = H // 8 * 8 W8 = W // 8 * 8 if flow_init is not None: flow_init_resize = F.interpolate( flow_init, size=(H8 // 8, W8 // 8), mode="nearest" ) flow_init_resize[:, :1] = ( flow_init_resize[:, :1].clone() * (W8 // 8 * 1.0) / flow_init.size()[3] ) flow_init_resize[:, 1:] = ( flow_init_resize[:, 1:].clone() * (H8 // 8 * 1.0) / flow_init.size()[2] ) if not hasattr(self.args, "not_use_rfr_mask") or ( hasattr(self.args, "not_use_rfr_mask") and (not self.args.not_use_rfr_mask) ): im18 = F.interpolate(image1, size=(H8 // 8, W8 // 8), mode="bilinear") im28 = F.interpolate(image2, size=(H8 // 8, W8 // 8), mode="bilinear") warp21 = backwarp(im28, flow_init_resize) error21 = torch.sum(torch.abs(warp21 - im18), dim=1, keepdim=True) # print('errormin', error21.min(), error21.max()) f12init = ( torch.exp( -self.attention2( torch.cat([im18, error21, flow_init_resize], dim=1) ) ** 2 ) * flow_init_resize ) else: flow_init_resize = None flow_init = torch.zeros( image1.size()[0], 2, image1.size()[2] // 8, image1.size()[3] // 8 ).cuda() error21 = torch.zeros( image1.size()[0], 1, image1.size()[2] // 8, image1.size()[3] // 8 ).cuda() f12_init = flow_init # print('None inital flow!') image1 = F.interpolate(image1, size=(H8, W8), mode="bilinear") image2 = F.interpolate(image2, size=(H8, W8), mode="bilinear") f12s, f12, f12_init = self.forward_pred( image1, image2, iters, flow_init_resize, upsample, test_mode ) if hasattr(self.args, "requires_sq_flow") and self.args.requires_sq_flow: for ii in range(len(f12s)): f12s[ii] = F.interpolate(f12s[ii], size=(H, W), mode="bilinear") f12s[ii][:, :1] = f12s[ii][:, :1].clone() / (1.0 * W8) * W f12s[ii][:, 1:] = f12s[ii][:, 1:].clone() / (1.0 * H8) * H if self.training: return f12s else: return [f12s[-1]], f12_init else: f12[:, :1] = f12[:, :1].clone() / (1.0 * W8) * W f12[:, 1:] = f12[:, 1:].clone() / (1.0 * H8) * H f12 = F.interpolate(f12, size=(H, W), mode="bilinear") # print('wo!!') return ( f12, f12_init, error21, ) def forward_pred( self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False ): """Estimate optical flow between pair of frames""" image1 = image1.contiguous() image2 = image2.contiguous() hdim = self.hidden_dim cdim = self.context_dim # run the feature network with autocast(device.type, enabled=self.args.mixed_precision): fmap1, fmap2 = self.fnet([image1, image2]) fmap1 = fmap1.float() fmap2 = fmap2.float() corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius) # run the context network with autocast(device.type, enabled=self.args.mixed_precision): cnet = self.fnet(image1) 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() if itr == 0: if flow_init is not None: coords1 = coords1 + flow_init corr = corr_fn(coords1) # index correlation volume flow = coords1 - coords0 with autocast(device.type, 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) return flow_predictions, flow_up, flow_init ####################### WARPING ####################### # expects batched tensors, considered low-level operation # img: bs, ch, h, w # flow: bs, xy (pix displace), h, w def flow_backwarp( img, flow, resample="bilinear", padding_mode="border", align_corners=False ): if len(img.shape) != 4: img = img[None,] if len(flow.shape) != 4: flow = flow[None,] q = ( 2 * flow / torch.tensor( [ flow.shape[-2], flow.shape[-1], ], device=flow.device, dtype=torch.float, )[None, :, None, None] ) q = q + torch.stack( torch.meshgrid( torch.linspace(-1, 1, flow.shape[-2]), torch.linspace(-1, 1, flow.shape[-1]), ) )[ None, ].to( flow.device ) if img.dtype != q.dtype: img = img.type(q.dtype) return nn.functional.grid_sample( img, q.flip(dims=(1,)).permute(0, 2, 3, 1), mode=resample, # nearest, bicubic, bilinear padding_mode=padding_mode, # border, zeros, reflection align_corners=align_corners, ) backwarp = flow_warp = flow_backwarp # mode: sum, avg, lin, softmax # lin/softmax w/out metric defaults to avg # must use gpu, move back to cpu if retain_device # typical metric: -20 * | img0 - backwarp(img1,flow) | # From Fannovel16: Changed mode params for common ops. def flow_forewarp( img, flow, mode="average", metric=None, mask=False, retain_device=True ): # setup #if mode == "sum": # mode = "summation" #elif mode == "avg": # mode = "average" if mode in ["lin", "linear"]: #mode = "linear" if metric is not None else "average" mode = "linear" if metric is not None else "avg" elif mode in ["sm", "softmax"]: #mode = "softmax" if metric is not None else "average" mode = "soft" if metric is not None else "avg" if len(img.shape) != 4: img = img[None,] if len(flow.shape) != 4: flow = flow[None,] if metric is not None and len(metric.shape) != 4: metric = metric[None,] flow = flow.flip(dims=(1,)) if img.dtype != torch.float32: img = img.type(torch.float32) if flow.dtype != torch.float32: flow = flow.type(torch.float32) if metric is not None and metric.dtype != torch.float32: metric = metric.type(torch.float32) # move to gpu if necessary assert img.device == flow.device if metric is not None: assert img.device == metric.device was_cpu = img.device.type == "cpu" if was_cpu: img = img.to("cuda") flow = flow.to("cuda") if metric is not None: metric = metric.to("cuda") # add mask if mask: bs, ch, h, w = img.shape img = torch.cat( [img, torch.ones(bs, 1, h, w, dtype=img.dtype, device=img.device)], dim=1 ) # forward, move back to cpu if desired ans = FunctionSoftsplat(img, flow, metric, mode) if was_cpu and retain_device: ans = ans.cpu() return ans forewarp = flow_forewarp # resizing utility def flow_resize(flow, size, mode="nearest", align_corners=False): # flow: bs,xy,h,w size = pixel_ij(size, rounding=True) if flow.dtype != torch.float: flow = flow.float() if len(flow.shape) == 3: flow = flow[None,] if flow.shape[-2:] == size: return flow return ( nn.functional.interpolate( flow, size=size, mode=mode, align_corners=align_corners if mode != "nearest" else None, ) * torch.tensor( [b / a for a, b in zip(flow.shape[-2:], size)], device=flow.device, )[None, :, None, None] ) ####################### TRADITIONAL ####################### # dense _lucaskanade = lambda a, b: np.moveaxis( cv2.optflow.calcOpticalFlowSparseToDense( a, b, # grid_step=5, sigma=0.5, ), 2, 0, )[ None, ] _farneback = lambda a, b: np.moveaxis( cv2.calcOpticalFlowFarneback( a, b, None, 0.6, 3, 25, 7, 5, 1.2, cv2.OPTFLOW_FARNEBACK_GAUSSIAN, ), 2, 0, )[ None, ] _dtvl1_ = cv2.optflow.createOptFlow_DualTVL1() _dtvl1 = lambda a, b: np.moveaxis( _dtvl1_.calc( a, b, None, ), 2, 0, )[ None, ] _simple = lambda a, b: np.moveaxis( cv2.optflow.calcOpticalFlowSF( a, b, 3, 5, 5, ), 2, 0, )[ None, ] _pca_ = cv2.optflow.createOptFlow_PCAFlow() _pca = lambda a, b: np.moveaxis( _pca_.calc( a, b, None, ), 2, 0, )[ None, ] _drlof = lambda a, b: np.moveaxis( cv2.optflow.calcOpticalFlowDenseRLOF( a, b, None, ), 2, 0, )[ None, ] _deepflow_ = cv2.optflow.createOptFlow_DeepFlow() _deepflow = lambda a, b: np.moveaxis( _deepflow_.calc( a, b, None, ), 2, 0, )[ None, ] def cv2flow(a, b, method="lucaskanade", back=False): if method == "lucaskanade": f = _lucaskanade a = a.convert("L").cv2() b = b.convert("L").cv2() elif method == "farneback": f = _farneback a = a.convert("L").cv2() b = b.convert("L").cv2() elif method == "dtvl1": f = _dtvl1 a = a.convert("L").cv2() b = b.convert("L").cv2() elif method == "simple": f = _simple a = a.convert("RGB").cv2() b = b.convert("RGB").cv2() elif method == "pca": f = _pca a = a.convert("L").cv2() b = b.convert("L").cv2() elif method == "drlof": f = _drlof a = a.convert("RGB").cv2() b = b.convert("RGB").cv2() elif method == "deepflow": f = _deepflow a = a.convert("L").cv2() b = b.convert("L").cv2() else: assert 0 ans = f(b, a) if back: ans = np.concatenate( [ ans, f(a, b), ] ) return torch.tensor(ans).flip(dims=(1,)) ####################### FLOWNET2 ####################### def flownet2(img_a, img_b, mode="shm", back=False): # package url = f"http://localhost:8109/get-flow" if mode == "shm": t = time.time() fn_a = img_a.save(mkfile(f"/dev/shm/_flownet2/{t}/img_a.png")) fn_b = img_b.save(mkfile(f"/dev/shm/_flownet2/{t}/img_b.png")) elif mode == "net": assert False, "not impl" q = u2d.img2uri(img.pil("RGB")) q.decode() resp = requests.get( url, params={ "img_a": fn_a, "img_b": fn_b, "mode": mode, "back": back, # 'vis': vis, }, ) # return ans = {"response": resp} if resp.status_code == 200: j = resp.json() ans["time"] = j["time"] ans["output"] = { "flow": torch.tensor(load(j["fn_flow"])), } # if vis: # ans['output']['vis'] = I(j['fn_vis']) if mode == "shm": shutil.rmtree(f"/dev/shm/_flownet2/{t}") return ans ####################### VISUALIZATION ####################### class Gridnet(nn.Module): def __init__(self, channels_0, channels_1, channels_2, total_dropout_p, depth): super().__init__() self.channels_0 = ch0 = channels_0 self.channels_1 = ch1 = channels_1 self.channels_2 = ch2 = channels_2 self.total_dropout_p = p = total_dropout_p self.depth = depth self.encoders = nn.ModuleList( [GridnetEncoder(ch0, ch1, ch2) for i in range(self.depth)] ) self.decoders = nn.ModuleList( [GridnetDecoder(ch0, ch1, ch2) for i in range(self.depth)] ) self.total_dropout = GridnetTotalDropout(p) return def forward(self, x): for e, enc in enumerate(self.encoders): t = [self.total_dropout(i) for i in t] if e != 0 else x t = enc(t) for d, dec in enumerate(self.decoders): t = [self.total_dropout(i) for i in t] t = dec(t) return t class GridnetEncoder(nn.Module): def __init__(self, channels_0, channels_1, channels_2): super().__init__() self.channels_0 = ch0 = channels_0 self.channels_1 = ch1 = channels_1 self.channels_2 = ch2 = channels_2 self.resnet_0 = GridnetResnet(ch0) self.resnet_1 = GridnetResnet(ch1) self.resnet_2 = GridnetResnet(ch2) self.downsample_01 = GridnetDownsample(ch0, ch1) self.downsample_12 = GridnetDownsample(ch1, ch2) return def forward(self, x): out = [ None, ] * 3 out[0] = self.resnet_0(x[0]) out[1] = self.resnet_1(x[1]) + self.downsample_01(out[0]) out[2] = self.resnet_2(x[2]) + self.downsample_12(out[1]) return out class GridnetDecoder(nn.Module): def __init__(self, channels_0, channels_1, channels_2): super().__init__() self.channels_0 = ch0 = channels_0 self.channels_1 = ch1 = channels_1 self.channels_2 = ch2 = channels_2 self.resnet_0 = GridnetResnet(ch0) self.resnet_1 = GridnetResnet(ch1) self.resnet_2 = GridnetResnet(ch2) self.upsample_10 = GridnetUpsample(ch1, ch0) self.upsample_21 = GridnetUpsample(ch2, ch1) return def forward(self, x): out = [ None, ] * 3 out[2] = self.resnet_2(x[2]) out[1] = self.resnet_1(x[1]) + self.upsample_21(out[2]) out[0] = self.resnet_0(x[0]) + self.upsample_10(out[1]) return out class GridnetConverter(nn.Module): def __init__(self, channels_in, channels_out): super().__init__() self.channels_in = cin = channels_in self.channels_out = cout = channels_out self.nets = nn.ModuleList( [ nn.Sequential( nn.PReLU(a), nn.Conv2d(a, b, kernel_size=1, padding=0), nn.BatchNorm2d(b), ) for a, b in zip(cin, cout) ] ) return def forward(self, x): return [m(q) for m, q in zip(self.nets, x)] class GridnetResnet(nn.Module): def __init__(self, channels): super().__init__() self.channels = ch = channels self.net = nn.Sequential( nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), ) return def forward(self, x): return x + self.net(x) class GridnetDownsample(nn.Module): def __init__(self, channels_in, channels_out): super().__init__() self.channels_in = chin = channels_in self.channels_out = chout = channels_out self.net = nn.Sequential( nn.PReLU(chin), nn.Conv2d(chin, chin, kernel_size=3, padding=1, stride=2), nn.BatchNorm2d(chin), nn.PReLU(chin), nn.Conv2d(chin, chout, kernel_size=3, padding=1), nn.BatchNorm2d(chout), ) return def forward(self, x): return self.net(x) class GridnetUpsample(nn.Module): def __init__(self, channels_in, channels_out): super().__init__() self.channels_in = chin = channels_in self.channels_out = chout = channels_out self.net = nn.Sequential( nn.Upsample(scale_factor=2, mode="nearest"), nn.PReLU(chin), nn.Conv2d(chin, chout, kernel_size=3, padding=1), nn.BatchNorm2d(chout), nn.PReLU(chout), nn.Conv2d(chout, chout, kernel_size=3, padding=1), nn.BatchNorm2d(chout), ) return def forward(self, x): return self.net(x) class GridnetTotalDropout(nn.Module): def __init__(self, p): super().__init__() assert 0 <= p < 1 self.p = p self.weight = 1 / (1 - p) return def get_drop(self, x): d = torch.rand(len(x))[:, None, None, None] < self.p d = (1 - d.float()).to(x.device) * self.weight return d def forward(self, x, force_drop=None): if force_drop is True: ans = x * self.get_drop(x) elif force_drop is False: ans = x else: if self.training: ans = x * self.get_drop(x) else: ans = x return ans class Interpolator(nn.Module): def __init__(self, size, mode="bilinear"): super().__init__() self.size = size self.mode = mode return def forward(self, x, is_flow=False): if x.shape[-2] == self.size: return x if len(x.shape) == 4: # bs,ch,h,w bs, ch, h, w = x.shape ans = nn.functional.interpolate( x, size=self.size, mode=self.mode, align_corners=(False, None)[self.mode == "nearest"], ) if is_flow: ans = ( ans * torch.tensor( [b / a for a, b in zip((h, w), self.size)], device=ans.device, )[None, :, None, None] ) return ans elif len(x.shape) == 5: # bs,k,ch,h,w (merge bs and k) bs, k, ch, h, w = x.shape return self.forward( x.view(bs * k, ch, h, w), is_flow=is_flow, ).view(bs, k, ch, *self.size) else: assert 0 ###################### CANNY ###################### def canny(img, a=100, b=200): img = I(img).convert("L") return I(cv2.Canny(img.cv2(), a, b)) # https://www.pyimagesearch.com/2015/04/06/zero-parameter-automatic-canny-edge-detection-with-python-and-opencv/ def canny_pis(img, sigma=0.33): # compute the median of the single channel pixel intensities img = I(img).convert("L").uint8(ch_last=False) v = np.median(img) # apply automatic Canny edge detection using the computed median lower = int(max(0, (1.0 - sigma) * v)) upper = int(min(255, (1.0 + sigma) * v)) edged = cv2.Canny(img[0], lower, upper) # return the edged image return I(edged) # https://en.wikipedia.org/wiki/Otsu%27s_method def canny_otsu(img): img = I(img).convert("L").uint8(ch_last=False) high, _ = cv2.threshold(img[0], 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) low = 0.5 * high return I(cv2.Canny(img[0], low, high)) def xdog(img, t=1.0, epsilon=0.04, phi=100, sigma=3, k=1.6): img = I(img).convert("L").uint8(ch_last=False) grey = np.asarray(img, dtype=np.float32) g0 = scipy.ndimage.gaussian_filter(grey, sigma) g1 = scipy.ndimage.gaussian_filter(grey, sigma * k) # ans = ((1+p) * g0 - p * g1) / 255 ans = (g0 - t * g1) / 255 ans = 1 + np.tanh(phi * (ans - epsilon)) * (ans < epsilon) return ans def dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True): img = I(img).convert("L").tensor()[None] kern0 = max(2 * int(sigma * kernel_factor) + 1, 3) kern1 = max(2 * int(sigma * k * kernel_factor) + 1, 3) g0 = kornia.filters.gaussian_blur2d( img, (kern0, kern0), (sigma, sigma), border_type="replicate", ) g1 = kornia.filters.gaussian_blur2d( img, (kern1, kern1), (sigma * k, sigma * k), border_type="replicate", ) ans = 0.5 + t * (g1 - g0) - epsilon ans = ans.clip(0, 1) if clip else ans return ans[0].numpy() # input: (bs,rgb(a),h,w) or (bs,1,h,w) # returns: (bs,1,h,w) def batch_dog(img, t=1.0, sigma=1.0, k=1.6, epsilon=0.01, kernel_factor=4, clip=True): # to grayscale if needed bs, ch, h, w = img.shape if ch in [3, 4]: img = kornia.color.rgb_to_grayscale(img[:, :3]) else: assert ch == 1 # calculate dog kern0 = max(2 * int(sigma * kernel_factor) + 1, 3) kern1 = max(2 * int(sigma * k * kernel_factor) + 1, 3) g0 = kornia.filters.gaussian_blur2d( img, (kern0, kern0), (sigma, sigma), border_type="replicate", ) g1 = kornia.filters.gaussian_blur2d( img, (kern1, kern1), (sigma * k, sigma * k), border_type="replicate", ) ans = 0.5 + t * (g1 - g0) - epsilon ans = ans.clip(0, 1) if clip else ans return ans ############### DERIVED DISTANCES ############### # input: (bs,h,w) or (bs,1,h,w) # returns: (bs,) # normalized s.t. metric is same across proportional image scales # average of two asymmetric distances # normalized by diameter and area def batch_chamfer_distance(gt, pred, block=1024, return_more=False): t = batch_chamfer_distance_t(gt, pred, block=block) p = batch_chamfer_distance_p(gt, pred, block=block) cd = (t + p) / 2 return cd def batch_chamfer_distance_t(gt, pred, block=1024, return_more=False): assert gt.device == pred.device and gt.shape == pred.shape bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] dpred = batch_edt(pred, block=block) cd = (gt * dpred).float().mean((-2, -1)) / np.sqrt(h**2 + w**2) if len(cd.shape) == 2: assert cd.shape[1] == 1 cd = cd.squeeze(1) return cd def batch_chamfer_distance_p(gt, pred, block=1024, return_more=False): assert gt.device == pred.device and gt.shape == pred.shape bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] dgt = batch_edt(gt, block=block) cd = (pred * dgt).float().mean((-2, -1)) / np.sqrt(h**2 + w**2) if len(cd.shape) == 2: assert cd.shape[1] == 1 cd = cd.squeeze(1) return cd # normalized by diameter # always between [0,1] def batch_hausdorff_distance(gt, pred, block=1024, return_more=False): assert gt.device == pred.device and gt.shape == pred.shape bs, h, w = gt.shape[0], gt.shape[-2], gt.shape[-1] dgt = batch_edt(gt, block=block) dpred = batch_edt(pred, block=block) hd = torch.stack( [ (dgt * pred).amax(dim=(-2, -1)), (dpred * gt).amax(dim=(-2, -1)), ] ).amax(dim=0).float() / np.sqrt(h**2 + w**2) if len(hd.shape) == 2: assert hd.shape[1] == 1 hd = hd.squeeze(1) return hd #################### UTILITIES #################### def reset_parameters(model): for layer in model.children(): if hasattr(layer, "reset_parameters"): layer.reset_parameters() return model def channel_squeeze(x, dim=1): a = x.shape[:dim] b = x.shape[dim + 2 :] return x.reshape(*a, -1, *b) def channel_unsqueeze(x, shape, dim=1): a = x.shape[:dim] b = x.shape[dim + 1 :] return x.reshape(*a, *shape, *b) def default_collate(items, device=None): return to(dict(torch.utils.data.dataloader.default_collate(items)), device) def to(x, device): if device is None: return x if issubclass(x.__class__, dict): return dict( { k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in x.items() } ) if isinstance(x, torch.Tensor): return x.to(device) if isinstance(x, np.ndarray): return torch.tensor(x).to(device) assert 0, "data not understood" ################ PARSING ################ from argparse import Namespace # args: all args # bargs: base args # pargs: data processing args # largs: data loading args # margs: model args # targs: training args # typically used to read dataset filters def read_filter(fn, cast=None, sort=True, sort_key=None): if cast is None: cast = lambda x: x ans = [cast(line) for line in read(fn).split("\n") if line != ""] if sort: return sorted(ans, key=sort_key) else: return ans ################ FILE MANAGEMENT ################ def mkfile(fn, parents=True, exist_ok=True): dn = "/".join(fn.split("/")[:-1]) mkdir(dn, parents=parents, exist_ok=exist_ok) return fn def mkdir(dn, parents=True, exist_ok=True): pathlib.Path(dn).mkdir(parents=parents, exist_ok=exist_ok) return dn if (not dn[-1] == "/" or dn == "/") else dn[:-1] def fstrip(fn, return_more=False): dspl = fn.split("/") dn = "/".join(dspl[:-1]) if len(dspl) > 1 else "." fn = dspl[-1] fspl = fn.split(".") if len(fspl) == 1: bn = fspl[0] ext = "" else: bn = ".".join(fspl[:-1]) ext = fspl[-1] if return_more: return Namespace( dn=dn, fn=fn, path=f"{dn}/{fn}", bn_path=f"{dn}/{bn}", bn=bn, ext=ext, ) else: return bn def read(fn, mode="r"): with open(fn, mode) as handle: return handle.read() def write(text, fn, mode="w"): mkfile(fn, parents=True, exist_ok=True) with open(fn, mode) as handle: return handle.write(text) import pickle def dump(obj, fn, mode="wb"): mkfile(fn, parents=True, exist_ok=True) with open(fn, mode) as handle: return pickle.dump(obj, handle) def load(fn, mode="rb"): with open(fn, mode) as handle: return pickle.load(handle) import json def jwrite(x, fn, mode="w", indent="\t", sort_keys=False): mkfile(fn, parents=True, exist_ok=True) with open(fn, mode) as handle: return json.dump(x, handle, indent=indent, sort_keys=sort_keys) def jread(fn, mode="r"): with open(fn, mode) as handle: return json.load(handle) try: import yaml def ywrite(x, fn, mode="w", default_flow_style=False): mkfile(fn, parents=True, exist_ok=True) with open(fn, mode) as handle: return yaml.dump(x, handle, default_flow_style=default_flow_style) def yread(fn, mode="r"): with open(fn, mode) as handle: return yaml.safe_load(handle) except: pass try: import pyunpack except: pass try: import mysql import mysql.connector except: pass ################ MISC ################ hakase = "./env/__hakase__.jpg" if not os.path.isfile(hakase): hakase = "./__env__/__hakase__.jpg" def mem(units="m"): return ( psProcess(os.getpid()).memory_info().rss / { "b": 1, "k": 1e3, "m": 1e6, "g": 1e9, "t": 1e12, }[units[0].lower()] ) def chunk(array, length, colwise=True): if colwise: return [array[i : i + length] for i in range(0, len(array), length)] else: return chunk(array, int(math.ceil(len(array) / length)), colwise=True) def classtree(x): return inspect.getclasstree(inspect.getmro(x)) ################ AESTHETIC ################ class Table: def __init__( self, table, delimiter=" ", orientation="br", double_colon=True, ): self.delimiter = delimiter self.orientation = orientation self.t = Table.parse(table, delimiter, orientation, double_colon) return # rendering def __str__(self): return self.render() def __repr__(self): return self.render() def render(self): # set up empty entry empty = ("", Table._spec(self.orientation, transpose=False)) # calculate table size t = copy.deepcopy(self.t) totalrows = len(t) totalcols = [len(r) for r in t] assert min(totalcols) == max(totalcols) totalcols = totalcols[0] # string-ify for i in range(totalrows): for j in range(totalcols): x, s = t[i][j] sp = s[11] if sp: x = eval(f'f"{{{x}{sp}}}"') Table._put((str(x), s), t, (i, j), empty) # expand delimiters _repl = ( lambda s: s[:2] + (1, 0, 0, 0, 0) + s[7:10] + (1,) + s[11:] if s[2] else s[:2] + (0, 0, 0, 0, 0) + s[7:10] + (1,) + s[11:] ) for i, row in enumerate(t): for j, (x, s_own) in enumerate(row): # expand delim_up(^) if s_own[3]: u, v = i, j while 0 <= u: _, s = t[u][v] if (i, j) != (u, v) and (s[2] and not s[10]): break Table._put((x, _repl(s)), t, (u, v), empty) u -= 1 # expand delim_down(v) if s_own[4]: u, v = i, j while u < totalrows: _, s = t[u][v] if (i, j) != (u, v) and (s[2] and not s[10]): break Table._put((x, _repl(s)), t, (u, v), empty) u += 1 # expand delim_right(>) if s_own[5]: u, v = i, j while v < totalcols: _, s = t[u][v] if (i, j) != (u, v) and (s[2] and not s[10]): break Table._put((x, _repl(s)), t, (u, v), empty) v += 1 # expand delim_left(<) if s_own[6]: u, v = i, j while 0 <= v: _, s = t[u][v] if (i, j) != (u, v) and (s[2] and not s[10]): break Table._put((x, _repl(s)), t, (u, v), empty) v -= 1 # justification calculation widths = [ 0, ] * totalcols # j heights = [ 0, ] * totalrows # i for i, row in enumerate(t): for j, (x, s) in enumerate(row): # height caclulation heights[i] = max(heights[i], x.count("\n")) # width calculation; non-delim fillers no contribution if s[2] or not s[10]: w = max(len(q) for q in x.split("\n")) widths[j] = max(widths[j], w) # no newline ==> height=1 heights = [h + 1 for h in heights] # render table rend = [] roff = 0 for i, row in enumerate(t): for j, (x, s) in enumerate(row): w, h = widths[j], heights[i] # expand fillers and delimiters if s[2] or s[10]: xs = x.split("\n") xw0 = min(len(l) for l in xs) xw1 = max(len(l) for l in xs) xh = len(xs) if (xw0 == xw1 == w) and (xh == h): pass elif xw0 == xw1 == w: x = "\n".join( [ xs[0], ] * h ) elif xh == h: x = "\n".join([(l[0] if l else "") * w for l in xs]) else: x = x[0] if x else " " x = "\n".join( [ x * w, ] * h ) # justify horizontally x = [l.rjust(w) if s[0] else l.ljust(w) for l in x.split("\n")] # justify vertically plus = [ " " * w, ] * (h - len(x)) x = plus + x if not s[1] else x + plus # input to table for r, xline in enumerate(x): Table._put(xline, rend, (roff + r, j), None) roff += h # return rendered string return "\n".join(["".join(r) for r in rend]) # parsing def _spec(s, transpose=False): if ":" in s: i = s.index(":") sp = s[i:] s = s[:i] else: sp = "" s = s.lower() return ( int("r" in s), # 0:: 0:left(l) 1:right(r) int("t" in s), # 1:: 0:bottom(b) 1:top(t) int(any([i in s for i in [".", "<", ">", "^", "v"]])), # 2:: delim_here(.) int("^" in s if not transpose else "<" in s), # 3:: delim_up(^) int("v" in s if not transpose else ">" in s), # 4:: delim_down(v) int(">" in s if not transpose else "v" in s), # 5:: delim_right(>) int("<" in s if not transpose else "^" in s), # 6:: delim_left(<) int("+" in s), # 7:: subtable(+) int("-" in s if not transpose else "|" in s), # 8:: subtable_horiz(-) int("|" in s if not transpose else "-" in s), # 9:: subtable_vert(|) int("_" in s), # 10:: fill(_); if delim, overwrite; else fit sp, # 11:: special(:) f-string for numbers ) def _put(obj, t, ij, empty): i, j = ij while i >= len(t): t.append([]) while j >= len(t[i]): t[i].append(empty) t[i][j] = obj return def parse( table, delimiter=" ", orientation="br", double_colon=True, ): # disabling transpose transpose = False # set up empty entry empty = ("", Table._spec(orientation, transpose)) # transpose t = [] for i, row in enumerate(table): for j, item in enumerate(row): ij = (i, j) if not transpose else (j, i) if type(item) == tuple and len(item) == 2 and type(item[1]) == str: item = (item[0], Table._spec(item[1], transpose)) elif double_colon and type(item) == str and "::" in item: x, s = item.split("::") item = (x, Table._spec(s, transpose)) else: item = (item, Table._spec(orientation, transpose)) Table._put(item, t, ij, empty) # normalization maxcol = 0 maxrow = len(t) for i, row in enumerate(t): # take element number into account maxcol = max(maxcol, len([i for i in row if not i[1][2]])) # take subtables into account for j, (x, s) in enumerate(row): if s[7]: r = len(x) maxrow = max(maxrow, i + r) c = max(len(q) for q in x) maxcol = max(maxcol, j + c) elif s[8]: c = len(x) maxcol = max(maxcol, j + c) elif s[9]: r = len(x) maxrow = max(maxrow, i + r) totalcols = 2 * maxcol + 1 totalrows = maxrow t += [[]] * (totalrows - len(t)) newt = [] delim = (delimiter, Table._spec("._" + orientation, transpose)) for i, row in enumerate(t): wasd = False tcount = 0 for j in range(totalcols): item = t[i][tcount] if tcount < len(t[i]) else empty isd = item[1][2] if wasd and isd: Table._put(empty, newt, (i, j), empty) wasd = False elif wasd and not isd: Table._put(item, newt, (i, j), empty) tcount += 1 wasd = False elif not wasd and isd: Table._put(item, newt, (i, j), empty) tcount += 1 wasd = True elif not wasd and not isd: Table._put(delim, newt, (i, j), empty) wasd = True t = newt # normalization: add dummy last column for delimiter for row in t: row.append(empty) # expand subtables delim_cols = [i for i in range(totalcols) if i % 2 == 0] while True: # find a table ij = None for i, row in enumerate(t): for j, item in enumerate(row): st, s = item if s[7]: ij = i, j, 7, st, s break elif s[8]: ij = i, j, 8, st, s break elif s[9]: ij = i, j, 9, st, s break if ij is not None: break if ij is None: break # replace its specs i, j, k, st, s = ij s = list(s) s[7] = s[8] = s[9] = 0 s = tuple(s) # expand it if k == 7: # 2d table for x, row in enumerate(st): for y, obj in enumerate(row): a = i + x if not transpose else i + y b = j + 2 * y if not transpose else j + 2 * x Table._put((obj, s), t, (a, b), None) if k == 8: # subtable_horiz for y, obj in enumerate(st): Table._put((obj, s), t, (i, j + 2 * y), None) if k == 9: # subtable_vert for x, obj in enumerate(st): Table._put((obj, s), t, (i + x, j), None) # return, finally return t class Resnet(nn.Module): def __init__(self, channels): super().__init__() self.channels = ch = channels self.net = nn.Sequential( nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), ) return def forward(self, x): return x + self.net(x) class Synthesizer(nn.Module): def __init__( self, size, channels_image, channels_flow, channels_mask, channels_feature ): super().__init__() self.size = size self.diam = diam(self.size) self.channels_image = cimg = channels_image self.channels_flow = cflow = channels_flow self.channels_mask = cmask = channels_mask self.channels_feature = cfeat = channels_feature self.channels = ch = cimg + cflow // 2 + cmask + cfeat self.interpolator = Interpolator(self.size, mode="bilinear") self.net = nn.Sequential( nn.Conv2d(ch + 3, 64, kernel_size=1, padding=0), Resnet(64), nn.Sequential( nn.PReLU(64), nn.Conv2d(64, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), ), Resnet(32), nn.Sequential( nn.PReLU(32), nn.Conv2d(32, 16, kernel_size=3, padding=1), nn.BatchNorm2d(16), ), Resnet(16), nn.Sequential( nn.PReLU(16), nn.Conv2d(16, 3, kernel_size=3, padding=1), ), ) return def forward(self, images, flows, masks, features, return_more=False): itp = self.interpolator images = [ (images[0] + images[1]) / 2, ] + images logimgs = [itp(pixel_logit(i[:, :3])) for i in images] cat = torch.cat( [ *logimgs, *[itp(f).norm(dim=1, keepdim=True) / self.diam for f in flows], *[itp(m) for m in masks], *[itp(f) for f in features], ], dim=1, ) residual = self.net(cat) return torch.sigmoid(logimgs[0] + 0.5 * residual), ( locals() if return_more else None ) class FlowZMetric(nn.Module): def __init__(self): super().__init__() return def forward(self, img0, img1, flow0, flow1, return_more=False): # B(i0,f0) = i1 # B(i1,f1) = i0 # F(x,f0,z0) # F(x,f1,z1) img0 = kornia.color.rgb_to_lab(img0[:, :3]) img1 = kornia.color.rgb_to_lab(img1[:, :3]) return [ -0.1 * (img1 - flow_backwarp(img0, flow0)).norm(dim=1, keepdim=True), # z0 -0.1 * (img0 - flow_backwarp(img1, flow1)).norm(dim=1, keepdim=True), # z1 ], (locals() if return_more else None) class NEDT(nn.Module): def __init__(self): super().__init__() return def forward( self, img, t=2.0, sigma_factor=1 / 540, k=1.6, epsilon=0.01, kernel_factor=4, exp_factor=540 / 15, return_more=False, ): with torch.no_grad(): dog = batch_dog( img, t=t, sigma=img.shape[-2] * sigma_factor, k=k, epsilon=epsilon, kernel_factor=kernel_factor, clip=False, ) edt = batch_edt((dog > 0.5).float()) ans = 1 - (-edt * exp_factor / max(edt.shape[-2:])).exp() return ans, (locals() if return_more else None) class HalfWarper(nn.Module): def __init__(self): super().__init__() self.channels_image = 4 * 3 self.channels_flow = 2 * 2 self.channels_mask = 2 * 1 self.channels = self.channels_image + self.channels_flow + self.channels_mask def morph_open(self, x, k): if k == 0: return x else: with torch.no_grad(): return kornia.morphology.opening(x, torch.ones(k, k, device=x.device)) def forward(self, img0, img1, flow0, flow1, z0, z1, k, t=0.5, return_more=False): # forewarps flow0_ = (1 - t) * flow0 flow1_ = t * flow1 f01 = forewarp(img0, flow1_, mode="sm", metric=z1, mask=True) f10 = forewarp(img1, flow0_, mode="sm", metric=z0, mask=True) f01i, f01m = f01[:, :-1], self.morph_open(f01[:, -1:], k=k) f10i, f10m = f10[:, :-1], self.morph_open(f10[:, -1:], k=k) # base guess base0 = f01m * f01i + (1 - f01m) * f10i base1 = f10m * f10i + (1 - f10m) * f01i ans = [ [ # images base0, base1, f01i, f10i, ], [ # flows flow0_, flow1_, ], [ # masks f01m, f10m, ], ] return ans, (locals() if return_more else None) class ResnetFeatureExtractor(nn.Module): def __init__(self, inferserve_query, size_in=None): super().__init__() self.inferserve_query = iq = inferserve_query self.size_in = si = size_in if iq[0] == "torchvision": # use pytorch pretrained resnet50 self.base_hparams = None resnet = tv.models.resnet50(pretrained=True) self.resize = T.Resize(256) self.resnet_preprocess = T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], ) self.conv1 = resnet.conv1 self.bn1 = resnet.bn1 self.relu = resnet.relu # 64ch, 128p (assuming 256p input) self.maxpool = resnet.maxpool self.layer1 = resnet.layer1 # 256ch, 64p self.layer2 = resnet.layer2 # 512ch, 32p else: base = userving.infer_model_load(*iq).eval() self.base_hparams = base.hparams self.resize = T.Resize(base.hparams.largs.size) self.resnet_preprocess = base.resnet_preprocess self.conv1 = base.resnet.conv1 self.bn1 = base.resnet.bn1 self.relu = base.resnet.relu # 64ch, 128p (assuming 256p input) self.maxpool = base.resnet.maxpool self.layer1 = base.resnet.layer1 # 256ch, 64p self.layer2 = base.resnet.layer2 # 512ch, 32p if self.size_in is None: self.sizes_out = None else: s = self.resize.size self.sizes_out = [ pixel_ij( rescale_dry(si, (s // 2) / si[0]), rounding="ceil" ), # conv1, 128p pixel_ij( rescale_dry(si, (s // 4) / si[0]), rounding="ceil" ), # layer1, 64p pixel_ij( rescale_dry(si, (s // 8) / si[0]), rounding="ceil" ), # layer2, 32p ] self.channels = [ 64, 256, 512, ] return def forward(self, x, force_sizes_out=False, return_more=False): ans = [] x = x[:, :3] x = self.resize(x) x = self.resnet_preprocess(x) x = self.conv1(x) x = self.bn1(x) x = self.relu(x) ans.append(x) # conv1 x = self.maxpool(x) x = self.layer1(x) ans.append(x) # layer1 x = self.layer2(x) ans.append(x) # layer2 if force_sizes_out or (self.sizes_out is None): self.sizes_out = [tuple(q.shape[-2:]) for q in ans] return ans, (locals() if return_more else None) class NetNedt(nn.Module): def __init__(self): super().__init__() chin = 3 + 1 + 4 + 4 + 1 + 1 ch = 16 chout = 1 self.net = nn.Sequential( nn.PReLU(chin), nn.Conv2d(chin, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, chout, kernel_size=3, padding=1), ) return def forward(self, out_base, out_base_nedt, hw_imgs, hw_masks, return_more=False): cat = torch.cat( [ out_base, # 3 out_base_nedt, # 1 hw_imgs[0], # 4 hw_imgs[1], # 4 hw_masks[0], # 1 hw_masks[1], # 1 ], dim=1, ) log = pixel_logit(cat.clip(0, 1)) ans = torch.sigmoid(self.net(log)) return ans, (locals() if return_more else None) class NetTail(nn.Module): def __init__(self): super().__init__() chin = 3 + 1 + 1 ch = 16 chout = 3 self.net = nn.Sequential( nn.PReLU(chin), nn.Conv2d(chin, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, ch, kernel_size=3, padding=1), nn.BatchNorm2d(ch), nn.PReLU(ch), nn.Conv2d(ch, chout, kernel_size=3, padding=1), ) return def forward(self, out_base, out_base_nedt, pred_nedt, return_more=False): cat = torch.cat( [ out_base, # 3 out_base_nedt, # 1 pred_nedt, # 1 ], dim=1, ) log = pixel_logit(cat.clip(0, 1)) ans = torch.sigmoid(log[:, :3] + self.net(log)) return ans, (locals() if return_more else None) class SoftsplatLite(nn.Module): def __init__(self): super().__init__() self.feature_extractor = ResnetFeatureExtractor( ("torchvision", "resnet50"), (540, 960), ) self.z_metric = FlowZMetric() self.flow_downsamplers = [ Interpolator(s, mode="bilinear") for s in self.feature_extractor.sizes_out ] self.gridnet_converter = GridnetConverter( self.feature_extractor.channels, [32, 64, 128], ) self.gridnet = Gridnet( *[32, 64, 128], total_dropout_p=0.0, depth=1, # equivalent to u-net ) self.nedt = NEDT() self.half_warper = HalfWarper() self.synthesizer = Synthesizer( (540, 960), self.half_warper.channels_image, self.half_warper.channels_flow, self.half_warper.channels_mask, self.gridnet.channels_0, ) return def forward(self, x, t=0.5, k=5, return_more=False): rm = return_more flow0, flow1 = x["flows"].swapaxes(0, 1) img0, img1 = x["images"][:, 0], x["images"][:, -1] (z0, z1), locs_z = self.z_metric(img0, img1, flow0, flow1, return_more=rm) img0 = torch.cat([img0, self.nedt(img0)[0]], dim=1) img1 = torch.cat([img1, self.nedt(img1)[0]], dim=1) # images and flows (hw_imgs, hw_flows, hw_masks), locs_hw = self.half_warper( img0, img1, flow0, flow1, z0, z1, k, t=t, return_more=rm, ) # features feats0, locs_fe0 = self.feature_extractor(img0, return_more=rm) feats1, locs_fe1 = self.feature_extractor(img1, return_more=rm) warps = [] for ft0, ft1, ds in zip(feats0, feats1, self.flow_downsamplers): (w, _, _), _ = self.half_warper( ft0, ft1, ds(flow0, 1), ds(flow1, 1), ds(z0), ds(z1), k, t=t, ) warps.append((w[0] + w[1]) / 2) feats = self.gridnet(self.gridnet_converter(warps)) # synthesis pred, locs_synth = self.synthesizer( hw_imgs, hw_flows, hw_masks, [ feats[0], ], return_more=rm, ) return pred, (locals() if rm else None) class DTM(nn.Module): def __init__(self): super().__init__() self.net_nedt = NetNedt() self.net_tail = NetTail() self.nedt = NEDT() return def forward(self, x, out_base, locs_base, return_more=False): rm = return_more with torch.no_grad(): out_base_nedt, locs_base_nedt = self.nedt(out_base, return_more=rm) hw_imgs, hw_masks = locs_base["hw_imgs"], locs_base["hw_masks"] pred_nedt, locs_nedt = self.net_nedt( out_base, out_base_nedt, hw_imgs, hw_masks, return_more=rm ) pred, locs_tail = self.net_tail( out_base, out_base_nedt, pred_nedt.clone().detach(), return_more=rm ) return torch.cat([pred, pred_nedt], dim=1), (locals() if rm else None) class RAFT(nn.Module): def __init__(self, path="/workspace/tensorrt/models/anime_interp_full.ckpt"): super().__init__() self.raft = RFR( Namespace( small=False, mixed_precision=False, ) ) if path is not None: sd = torch.load(path)["model_state_dict"] self.raft.load_state_dict( { k[len("module.flownet.") :]: v for k, v in sd.items() if k.startswith("module.flownet.") }, strict=False, ) return def forward(self, img0, img1, flow0=None, iters=12, return_more=False): if flow0 is not None: flow0 = flow0.flip(dims=(1,)) out = self.raft(img1, img0, iters=iters, flow_init=flow0) return out[0].flip(dims=(1,)), (locals() if return_more else None)