Merge pull request #8 from FuouM/stab

StabStitch

README.md

@@ -9,6 +9,8 @@ Click name to jump to workflow
 4. [**CoLIE**](#colie). Paper: [Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations](https://github.com/ctom2/colie)
 5. [**VFIMamba**](#vfimamba). Paper: [Video Frame Interpolation with State Space Models](https://github.com/MCG-NJU/VFIMamba)
 6. [**ConvIR**](#convir). Paper: [Revitalizing Convolutional Network for Image Restoration](https://github.com/c-yn/ConvIR)
+7. [**StabStitch**](#stabstitch). Paper: [Eliminating Warping Shakes for Unsupervised Online Video Stitching](https://github.com/nie-lang/StabStitch)
+
 
 ## Workflows
 
@@ -26,7 +28,7 @@ Download the model [here](https://github.com/Yuehan717/RealViformer?tab=readme-o
 
 [workflow_realviformer.json](workflow/workflow_realviformer.json)
 
-![realviformer_example](example_realviformer.png)
+![realviformer_example](workflow_examples/example_realviformer.png)
 
 (Not a workflow-embedded image)
 
@@ -75,7 +77,7 @@ To install, `pip install [].whl`
 
 [workflow_vfi_mamba.json](workflow/workflow_vfi_mamba.json)
 
-![example_vfi_mamba](example_vfi_mamba.png)
+![example_vfi_mamba](workflow_examples/example_vfi_mamba.png)
 
 (Not a workflow-embedded image)
 
@@ -134,6 +136,24 @@ models\convir
     convir_rsblur.pkl
 ```
 
+### StabStitch
+
+Download all 3 models in the [Code - Pre-trained model](https://github.com/nie-lang/StabStitch?tab=readme-ov-file#pre-trained-model) section.
+
+```
+models/stabstitch/temporal_warp.pth
+models/stabstitch/spatial_warp.pth
+models/stabstitch/smooth_warp.pth
+```
+
+Use `interpolate_mode = NORMAL` or `do_linear_blend = True` to eliminate dark borders. Inputs will be resized to 360x480. Recommends using `StabStitch Crop Resize`.
+
+| StabStitch | StabStitch Stabilize |
+|-|-|
+| [stabstitch_stitch.json](workflow/stabstitch_stitch.json) (Example videos in `examples\stabstitch`) | [stabstich_stabilize.json](workflow/stabstitch_stabilize.json) |
+| ![example_stabstitch_stitch](workflow_examples/example_stabstitch_stitch.png) | ![example_stabstitch_stabilize](workflow_examples/example_stabstitch_stabilize.png) | 
+
+(Not workflow-embedded images)
 
 ## Credits
 
@@ -178,4 +198,13 @@ models\convir
   year={2023},
   organization={PMLR}
 }
-```
+```
+
+```
+@article{nie2024eliminating,
+  title={Eliminating Warping Shakes for Unsupervised Online Video Stitching},
+  author={Nie, Lang and Lin, Chunyu and Liao, Kang and Zhang, Yun and Liu, Shuaicheng and Zhao, Yao},
+  journal={arXiv preprint arXiv:2403.06378},
+  year={2024}
+}
+```

__init__.py

@@ -1,11 +1,20 @@
-from .run import WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance
-
+from .run import (
+    WFEN,
+    CoLIE_LowLight_Enhance,
+    ProPIH_Harmonizer,
+    RealViFormerSR,
+)
 from .run_convir import (
+    ConvIR_DefocusDeblur,
     ConvIR_DeHaze,
     ConvIR_DeRain,
     ConvIR_DeSnow,
     ConvIR_MotionDeBlur,
-    ConvIR_DefocusDeblur,
+)
+from .run_stabstitch import (
+    StabStitch_Crop_Resize,
+    StabStitch_Stabilize,
+    StabStitch_Stitch,
 )
 
 NODE_CLASS_MAPPINGS = {
@@ -18,6 +27,9 @@ NODE_CLASS_MAPPINGS = {
     "ConvIR_DeSnow": ConvIR_DeSnow,
     "ConvIR_MotionDeBlur": ConvIR_MotionDeBlur,
     "ConvIR_DefocusDeblur": ConvIR_DefocusDeblur,
+    "StabStitch": StabStitch_Stitch,
+    "StabStitch_Stabilize": StabStitch_Stabilize,
+    "StabStitch_Crop_Resize": StabStitch_Crop_Resize,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -30,6 +42,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ConvIR_DeSnow": "ConvIR DeSnow",
     "ConvIR_MotionDeBlur": "ConvIR Motion DeBlur",
     "ConvIR_DefocusDeblur": "ConvIR Defocus Deblur",
+    "StabStitch": "StabStitch",
+    "StabStitch_Stabilize": "StabStitch Stabilize",
+    "StabStitch_Crop_Resize": "StabStitch Crop Resize",
 }
 
 try:

constants.py

@@ -8,3 +8,9 @@ VFI_MAMBA_PATHS = {
     "VFIMamba_S": "models/vfimamba/VFIMamba_S.pkl",
 }
 VFI_MAMBA_DEFAULT = "VFIMamba"
+
+STAB_TEMPORAL_PATH = "models/stabstitch/temporal_warp.pth"
+STAB_SPATIAL_PATH = "models/stabstitch/spatial_warp.pth"
+STAB_SMOOTH_PATH = "models/stabstitch/smooth_warp.pth"
+STAB_CONSTRUCT_MODES = ["FAST", "NORMAL"]
+STAB_CONSTRUCT_DEFAULT = "FAST"

module_stabstitch/smooth_network.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+
+grid_h = 6
+grid_w = 8
+
+
+def build_SmoothNet(net, tsmotion_list, smesh_list, omask_tensor_list):
+    # predict the delta motion for tsflow
+    stitch_mesh, overlap_mask, ori_path, delta_motion = net(
+        smesh_list, omask_tensor_list, tsmotion_list
+    )
+
+    # get the smoothes tsflow
+    smooth_path = ori_path + delta_motion
+    # get the actual warping mesh
+    target_mesh = stitch_mesh - delta_motion  # bs, T, h, w, 2
+
+    out_dict = {}
+    out_dict.update(
+        ori_path=ori_path,
+        smooth_path=smooth_path,
+        ori_mesh=stitch_mesh,
+        target_mesh=target_mesh,
+        overlap_mask=overlap_mask,
+        delta_motion=delta_motion,
+    )
+
+    return out_dict
+
+
+# define and forward ( Because of the load is unbalanced when use torch.nn.DataParallel, we define warp in forward)
+class SmoothNet(nn.Module):
+    def __init__(self, dropout=0.0):
+        super(SmoothNet, self).__init__()
+
+        self.MotionPre = MotionPre1()
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    # forward
+    def forward(self, smesh_list, omask_list, tsmotion_list):
+        # to generate meshflow from the first motion
+        tsflow_list = [tsmotion_list[0]]
+        for i in range(1, len(tsmotion_list)):
+            tsflow_list.append(tsflow_list[i - 1] + tsmotion_list[i])
+
+        # generate meshflow tensor
+        smesh = torch.cat(smesh_list, 3)  # bs, h, w, 2*T
+        omask = torch.cat(omask_list, 3)  # bs, h, w, 2*T
+        tsflow = torch.cat(tsflow_list, 3)  # bs, h, w, 2*T
+
+        smesh = smesh.reshape(
+            -1, grid_h + 1, grid_w + 1, len(smesh_list), 2
+        )  # bs, h, w, T, 2
+        smesh = smesh.permute(0, 3, 1, 2, 4)  # bs, T, h, w, 2
+
+        omask = omask.reshape(
+            -1, grid_h + 1, grid_w + 1, len(omask_list), 1
+        )  # bs, h, w, T, 2
+        omask = omask.permute(0, 3, 1, 2, 4)  # bs, T, h, w, 2
+
+        tsflow = tsflow.reshape(
+            -1, grid_h + 1, grid_w + 1, len(tsflow_list), 2
+        )  # bs, h, w, T, 2
+        tsflow = tsflow.permute(0, 3, 1, 2, 4)  # bs, T, h, w, 2
+
+        delta_tsflow = self.MotionPre(smesh, omask, tsflow)
+
+        return smesh, omask, tsflow, delta_tsflow
+
+
+class MotionPre1(nn.Module):
+    def __init__(self, kernel=5):
+        super().__init__()
+
+        self.embedding1 = nn.Sequential(nn.Linear(2, 24), nn.ReLU())
+        self.embedding2 = nn.Sequential(nn.Linear(1, 8), nn.ReLU())
+        self.embedding3 = nn.Sequential(nn.Linear(2, 32), nn.ReLU())
+
+        self.pad = kernel // 2
+        self.MotionConv3D = nn.Sequential(
+            nn.Conv3d(64, 64, (kernel, 3, 3), padding=(self.pad, 1, 1)),
+            nn.ReLU(),
+            nn.Conv3d(64, 64, (kernel, 3, 3), padding=(self.pad, 1, 1)),
+            nn.ReLU(),
+            nn.Conv3d(64, 64, (kernel, 3, 3), padding=(self.pad, 1, 1)),
+            nn.ReLU(),
+        )
+
+        self.decoding = nn.Sequential(nn.Linear(64, 2))
+
+    def forward(self, smesh, omask, tsflow):
+        # input: meshflow -- bs, T, h, w, 2
+        # output: delta_meshflow -- bs, T, h, w, 2
+
+        # predict the delta motion to stabilize the last meshflow
+        hidden1 = self.embedding1(smesh)  # bs, T, H, W, 32
+        hidden2 = self.embedding2(omask)  # bs, T, H, W, 32
+        hidden3 = self.embedding3(tsflow)  # bs, T, H, W, 32
+
+        hidden = torch.cat([hidden1, hidden2, hidden3], 4)  # bs, T, H, W, 64
+        hidden = self.MotionConv3D(hidden.permute(0, 4, 1, 2, 3))  # bs, 64, T, H, W
+        delta_tsflow = self.decoding(hidden.permute(0, 2, 3, 4, 1))  # bs, T, H, W, 2
+
+        return delta_tsflow

module_stabstitch/spatial_network.py

@@ -0,0 +1,382 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+from . import torch_DLT, torch_homo_transform
+
+grid_h = 6
+grid_w = 8
+
+
+# Covert global homo into mesh
+def H2Mesh(H, rigid_mesh):
+    H_inv = torch.inverse(H)
+    ori_pt = rigid_mesh.reshape(rigid_mesh.size()[0], -1, 2)
+    ones = torch.ones(rigid_mesh.size()[0], (grid_h + 1) * (grid_w + 1), 1)
+    if torch.cuda.is_available():
+        ori_pt = ori_pt.cuda()
+        ones = ones.cuda()
+
+    ori_pt = torch.cat((ori_pt, ones), 2)  # bs*(grid_h+1)*(grid_w+1)*3
+    tar_pt = torch.matmul(H_inv, ori_pt.permute(0, 2, 1))  # bs*3*(grid_h+1)*(grid_w+1)
+
+    mesh_x = torch.unsqueeze(tar_pt[:, 0, :] / tar_pt[:, 2, :], 2)
+    mesh_y = torch.unsqueeze(tar_pt[:, 1, :] / tar_pt[:, 2, :], 2)
+    mesh = torch.cat((mesh_x, mesh_y), 2).reshape(
+        [rigid_mesh.size()[0], grid_h + 1, grid_w + 1, 2]
+    )
+
+    return mesh
+
+
+# get rigid mesh
+def get_rigid_mesh(batch_size, height, width):
+    ww = torch.matmul(
+        torch.ones([grid_h + 1, 1]),
+        torch.unsqueeze(torch.linspace(0.0, float(width), grid_w + 1), 0),
+    )
+    hh = torch.matmul(
+        torch.unsqueeze(torch.linspace(0.0, float(height), grid_h + 1), 1),
+        torch.ones([1, grid_w + 1]),
+    )
+    if torch.cuda.is_available():
+        ww = ww.cuda()
+        hh = hh.cuda()
+
+    ori_pt = torch.cat((ww.unsqueeze(2), hh.unsqueeze(2)), 2)  # (grid_h+1)*(grid_w+1)*2
+    ori_pt = ori_pt.unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+    return ori_pt
+
+
+# normalize mesh from -1 ~ 1
+def get_norm_mesh(mesh, height, width):
+    batch_size = mesh.size()[0]
+    mesh_w = mesh[..., 0] * 2.0 / float(width) - 1.0
+    mesh_h = mesh[..., 1] * 2.0 / float(height) - 1.0
+    norm_mesh = torch.stack([mesh_w, mesh_h], 3)  # bs*(grid_h+1)*(grid_w+1)*2
+
+    return norm_mesh.reshape([batch_size, -1, 2])  # bs*-1*2
+
+
+def build_SpatialNet(net, input1_tensor, input2_tensor):
+    batch_size, _, img_h, img_w = input1_tensor.size()
+
+    # input1_tensor = (input1_tensor / 127.5) - 1.0
+    # input2_tensor = (input2_tensor / 127.5) - 1.0
+
+    H_motion, mesh_motion = net(input1_tensor, input2_tensor)
+
+    H_motion = H_motion.reshape(-1, 4, 2)
+    mesh_motion = mesh_motion.reshape(-1, grid_h + 1, grid_w + 1, 2)
+
+    # initialize the source points bs x 4 x 2
+    src_p = torch.tensor([[0.0, 0.0], [img_w, 0.0], [0.0, img_h], [img_w, img_h]])
+    if torch.cuda.is_available():
+        src_p = src_p.cuda()
+    src_p = src_p.unsqueeze(0).expand(batch_size, -1, -1)
+    # target points
+    dst_p = src_p + H_motion
+    # solve homo using DLT
+    H = torch_DLT.tensor_DLT(src_p, dst_p)
+
+    rigid_mesh = get_rigid_mesh(batch_size, img_h, img_w)
+    ini_mesh = H2Mesh(H, rigid_mesh)
+    mesh = ini_mesh + mesh_motion
+
+    # calculate the overlapping regions to apply shape-preserving constraints
+    overlap_one = torch.ones_like(mesh[..., 0])
+    overlap_zero = torch.zeros_like(mesh[..., 0])
+
+    overlap_w = torch.where(
+        (mesh[..., 0] < img_w) & (mesh[..., 0] >= 0), overlap_one, overlap_zero
+    )
+    overlap_h = torch.where(
+        (mesh[..., 1] < img_h) & (mesh[..., 1] >= 0), overlap_one, overlap_zero
+    )
+    overlap_mesh = overlap_w * overlap_h
+    overlap_mesh = torch.clamp(overlap_mesh.unsqueeze(3), 0, 1)
+
+    out_dict = {}
+    out_dict.update(motion=mesh - rigid_mesh, overlap_mesh=overlap_mesh)
+
+    return out_dict
+
+
+def get_res18_FeatureMap(resnet18_model):
+    # feature extractor
+    # suppose input resolution to be 256*256   |   512*384
+    # the output resolution should be 16*16    |    32*24
+
+    layers_list = []
+
+    # mask
+    # layers_list.append(nn.Conv2d(6, 3, kernel_size=3, padding=1, bias=False))
+
+    layers_list.append(resnet18_model.conv1)  # stride 2*2     H/2
+    # layers_list.append(nn.Conv2d(6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False))
+    layers_list.append(resnet18_model.bn1)
+    layers_list.append(resnet18_model.relu)
+    layers_list.append(resnet18_model.maxpool)  # stride 2       H/4
+
+    layers_list.append(resnet18_model.layer1)  # H/4
+    layers_list.append(resnet18_model.layer2)  # H/8
+
+    feature_extractor_stage1 = nn.Sequential(*layers_list)
+    feature_extractor_stage2 = nn.Sequential(resnet18_model.layer3)
+
+    # for img of 360*640
+    # the resolution of feature1 should be 45*80
+    # the resolution of feature1 should be *40
+
+    return feature_extractor_stage1, feature_extractor_stage2
+
+
+# define and forward
+class SpatialNet(nn.Module):
+    def __init__(self):
+        super(SpatialNet, self).__init__()
+
+        self.regressNet1_part1 = nn.Sequential(
+            nn.Conv2d(2, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 12, 20
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 6, 10
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 3, 5
+        )
+
+        self.regressNet1_part2 = nn.Sequential(
+            nn.Linear(in_features=768, out_features=512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=512, out_features=128, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=128, out_features=8, bias=True),
+        )
+
+        self.regressNet2_part1 = nn.Sequential(
+            nn.Conv2d(121, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 23, 40
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 12, 20
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 6, 10
+            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 3, 5
+        )
+
+        self.regressNet2_part2 = nn.Sequential(
+            nn.Linear(in_features=1536, out_features=1024, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=1024, out_features=512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(
+                in_features=512, out_features=(grid_w + 1) * (grid_h + 1) * 2, bias=True
+            ),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+        resnet18_model = models.resnet.resnet18(weights="DEFAULT")
+        if torch.cuda.is_available():
+            resnet18_model = resnet18_model.cuda()
+        self.feature_extractor_stage1, self.feature_extractor_stage2 = (
+            get_res18_FeatureMap(resnet18_model)
+        )
+        # -----------------------------------------
+
+    # forward
+    def forward(self, input1_tesnor, input2_tesnor):
+        batch_size, _, img_h, img_w = input1_tesnor.size()
+
+        # print("111")
+        feature_1_64 = self.feature_extractor_stage1(input1_tesnor)
+        feature_1_32 = self.feature_extractor_stage2(feature_1_64)
+        feature_2_64 = self.feature_extractor_stage1(input2_tesnor)
+        feature_2_32 = self.feature_extractor_stage2(feature_2_64)
+        # print("222")
+        ######### stage 1
+        correlation_32 = self.CCL(feature_1_32, feature_2_32)
+        temp_1 = self.regressNet1_part1(correlation_32)
+        temp_1 = temp_1.view(temp_1.size()[0], -1)
+        offset_1 = self.regressNet1_part2(temp_1)
+        H_motion_1 = offset_1.reshape(-1, 4, 2)
+
+        # print("333")
+        src_p = torch.tensor(
+            [[0.0, 0.0], [img_w - 1, 0.0], [0.0, img_h - 1], [img_w - 1, img_h - 1]]
+        )
+        if torch.cuda.is_available():
+            src_p = src_p.cuda()
+        src_p = src_p.unsqueeze(0).expand(batch_size, -1, -1)
+        dst_p = src_p + H_motion_1
+        H = torch_DLT.tensor_DLT(src_p / 8, dst_p / 8)
+
+        M_tensor = torch.tensor(
+            [
+                [img_w / 8 / 2.0, 0.0, img_w / 8 / 2.0],
+                [0.0, img_h / 8 / 2.0, img_h / 8 / 2.0],
+                [0.0, 0.0, 1.0],
+            ]
+        )
+
+        if torch.cuda.is_available():
+            M_tensor = M_tensor.cuda()
+
+        M_tile = M_tensor.unsqueeze(0).expand(batch_size, -1, -1)
+        M_tensor_inv = torch.inverse(M_tensor)
+        M_tile_inv = M_tensor_inv.unsqueeze(0).expand(batch_size, -1, -1)
+        H_mat = torch.matmul(torch.matmul(M_tile_inv, H), M_tile)
+
+        warp_feature_2_64 = torch_homo_transform.transformer(
+            feature_2_64, H_mat, (int(img_h / 8), int(img_w / 8))
+        )
+        # print("444")
+        ######### stage 2
+        # correlation_64 = self.CCL(feature_1_64, warp_feature_2_64)
+        correlation_64 = self.cost_volume(
+            feature_1_64, warp_feature_2_64, search_range=5, norm=False
+        )  # bs, 9^2, 45, 80
+        temp_2 = self.regressNet2_part1(correlation_64)
+        temp_2 = temp_2.view(temp_2.size()[0], -1)
+        offset_2 = self.regressNet2_part2(temp_2)
+        # print("555")
+
+        return offset_1, offset_2
+
+    @staticmethod
+    def cost_volume(x1, x2, search_range, norm=True, fast=True):
+        if norm:
+            x1 = F.normalize(x1, p=2, dim=1)
+            x2 = F.normalize(x2, p=2, dim=1)
+        bs, c, h, w = x1.shape
+        padded_x2 = F.pad(x2, [search_range] * 4)  # [b,c,h,w] -> [b,c,h+sr*2,w+sr*2]
+        max_offset = search_range * 2 + 1
+
+        if fast:
+            # faster(*2) but cost higher(*n) GPU memory
+            patches = F.unfold(padded_x2, (max_offset, max_offset)).reshape(
+                bs, c, max_offset**2, h, w
+            )
+            cost_vol = (x1.unsqueeze(2) * patches).mean(dim=1, keepdim=False)
+        else:
+            # slower but save memory
+            cost_vol = []
+            for j in range(0, max_offset):
+                for i in range(0, max_offset):
+                    x2_slice = padded_x2[:, :, j : j + h, i : i + w]
+                    cost = torch.mean(x1 * x2_slice, dim=1, keepdim=True)
+                    cost_vol.append(cost)
+            cost_vol = torch.cat(cost_vol, dim=1)
+
+        cost_vol = F.leaky_relu(cost_vol, 0.1)
+
+        return cost_vol
+
+    def extract_patches(self, x, kernel=3, stride=1):
+        if kernel != 1:
+            x = nn.ZeroPad2d(1)(x)
+        x = x.permute(0, 2, 3, 1)
+        all_patches = x.unfold(1, kernel, stride).unfold(2, kernel, stride)
+        return all_patches
+
+    def CCL(self, feature_1, feature_2):
+        bs, c, h, w = feature_1.size()
+
+        norm_feature_1 = F.normalize(feature_1, p=2, dim=1)
+        norm_feature_2 = F.normalize(feature_2, p=2, dim=1)
+        # print(norm_feature_2.size())
+
+        patches = self.extract_patches(norm_feature_2)
+        if torch.cuda.is_available():
+            patches = patches.cuda()
+
+        matching_filters = patches.reshape(
+            (
+                patches.size()[0],
+                -1,
+                patches.size()[3],
+                patches.size()[4],
+                patches.size()[5],
+            )
+        )
+
+        match_vol = []
+        for i in range(bs):
+            single_match = F.conv2d(
+                norm_feature_1[i].unsqueeze(0), matching_filters[i], padding=1
+            )
+            match_vol.append(single_match)
+
+        match_vol = torch.cat(match_vol, 0)
+        # print(match_vol .size())
+
+        # scale softmax
+        softmax_scale = 10
+        match_vol = F.softmax(match_vol * softmax_scale, 1)
+
+        channel = match_vol.size()[1]
+
+        h_one = torch.linspace(0, h - 1, h)
+        one1w = torch.ones(1, w)
+        if torch.cuda.is_available():
+            h_one = h_one.cuda()
+            one1w = one1w.cuda()
+        h_one = torch.matmul(h_one.unsqueeze(1), one1w)
+        h_one = h_one.unsqueeze(0).unsqueeze(0).expand(bs, channel, -1, -1)
+
+        w_one = torch.linspace(0, w - 1, w)
+        oneh1 = torch.ones(h, 1)
+        if torch.cuda.is_available():
+            w_one = w_one.cuda()
+            oneh1 = oneh1.cuda()
+        w_one = torch.matmul(oneh1, w_one.unsqueeze(0))
+        w_one = w_one.unsqueeze(0).unsqueeze(0).expand(bs, channel, -1, -1)
+
+        c_one = torch.linspace(0, channel - 1, channel)
+        if torch.cuda.is_available():
+            c_one = c_one.cuda()
+        c_one = c_one.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand(bs, -1, h, w)
+
+        flow_h = match_vol * (c_one // w - h_one)
+        flow_h = torch.sum(flow_h, dim=1, keepdim=True)
+        flow_w = match_vol * (c_one % w - w_one)
+        flow_w = torch.sum(flow_w, dim=1, keepdim=True)
+
+        feature_flow = torch.cat([flow_w, flow_h], 1)
+
+        return feature_flow

module_stabstitch/stabstitch_model.py

@@ -0,0 +1,525 @@
+import torch
+import torch.nn.functional as F
+import tqdm
+from torchvision.transforms import GaussianBlur
+
+from . import torch_tps_transform, torch_tps_transform_point
+from .smooth_network import SmoothNet, build_SmoothNet
+from .spatial_network import SpatialNet, build_SpatialNet
+from .temporal_network import TemporalNet, build_TemporalNet
+
+grid_h = 6
+grid_w = 8
+
+img_h = 360
+img_w = 480
+
+
+def crop_or_resize(tensor, target_h, target_w, crop_side="left"):
+    B, H, W, C = tensor.shape
+    aspect_ratio = W / H
+    target_aspect_ratio = target_w / target_h
+
+    if aspect_ratio == target_aspect_ratio:
+        # Just resize
+        resized = F.interpolate(
+            tensor.permute(0, 3, 1, 2),
+            size=(target_h, target_w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return resized.permute(0, 2, 3, 1)
+
+    elif aspect_ratio > target_aspect_ratio:
+        # Crop width, then resize
+        new_w = int(H * target_aspect_ratio)
+        if crop_side == "left":
+            cropped = tensor[:, :, :new_w, :]
+        elif crop_side == "right":
+            cropped = tensor[:, :, -new_w:, :]
+        else:
+            raise ValueError("crop_side must be 'left' or 'right'")
+
+        resized = F.interpolate(
+            cropped.permute(0, 3, 1, 2),
+            size=(target_h, target_w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return resized.permute(0, 2, 3, 1)
+
+    else:
+        # Crop height, then resize
+        new_h = int(W / target_aspect_ratio)
+        cropped = tensor[:, :new_h, :, :]
+
+        resized = F.interpolate(
+            cropped.permute(0, 3, 1, 2),
+            size=(target_h, target_w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return resized.permute(0, 2, 3, 1)
+
+
+def video_preprocess(src_video: torch.Tensor, h: int, w: int):
+    # [B, H, W, C] [0, 1]
+    tensor = src_video.permute(0, 3, 1, 2)  # [B, C, H, W]
+    tensor = tensor[:, [2, 1, 0], :, :]  # RGB -> BGR
+    tensor = tensor * 2 - 1  # [-1, 1]
+    tensor = F.interpolate(tensor, size=(h, w), mode="bilinear", align_corners=False)
+    return tensor
+
+
+def video_to_list_ts(video: torch.Tensor):
+    video_list = [tensor.unsqueeze(0) for tensor in video]
+    return video_list
+
+
+def stabstitch_inference(
+    src_video_1: torch.Tensor,
+    src_video_2: torch.Tensor,
+    spatial_path: str,
+    temporal_path: str,
+    smooth_path: str,
+    device,
+    mode: str,
+    do_linear_blend: bool,
+):
+    spatial_net, temporal_net, smooth_net = load_models(
+        spatial_path, temporal_path, smooth_path, device
+    )
+
+    video_1 = video_preprocess(src_video_1, img_h, img_w).to(device)
+    video_2 = video_preprocess(src_video_2, img_h, img_w).to(device)
+
+    video_list_1 = video_to_list_ts(video_1)
+    video_list_2 = video_to_list_ts(video_2)
+
+    return inference(
+        mode,
+        do_linear_blend,
+        spatial_net,
+        temporal_net,
+        smooth_net,
+        video_list_1,
+        video_list_2,
+    )
+
+
+def inference(
+    mode,
+    do_linear_blend,
+    spatial_net,
+    temporal_net,
+    smooth_net,
+    video_list_1,
+    video_list_2,
+):
+    num_frames = len(video_list_2)
+    # motion estimation
+    smotion_tensor_list, omask_tensor_list = spatial_warping(
+        spatial_net, video_list_1, video_list_2, num_frames
+    )
+
+    # step 2: temporal warp
+    with torch.no_grad():
+        temporal_batch_out = build_TemporalNet(temporal_net, video_list_2)
+    tmotion_tensor_list = temporal_batch_out["motion_list"]
+
+    ##############################################
+    #############   data preparation  ############
+    # converting tmotion (t-th frame) into tsmotion ( (t-1)-th frame )
+
+    smesh_list, tsmotion_list = temporaling(smotion_tensor_list, tmotion_tensor_list)
+
+    # step 3: smooth warp
+    target_mesh = smooth_meshing(
+        smooth_net,
+        smotion_tensor_list,
+        omask_tensor_list,
+        tmotion_tensor_list,
+        smesh_list,
+        tsmotion_list,
+    )
+
+    stable_list, mask_list_2, img_2_list, _, _ = get_stable_sqe(
+        video_list_1, video_list_2, target_mesh, mode, do_linear_blend
+    )
+
+    stable_list = [stable.unsqueeze(0)[:, [2, 1, 0], :, :] for stable in stable_list]
+
+    return stable_list, img_2_list, mask_list_2
+
+
+def stabstitch_single_inference(
+    src_video_1: torch.Tensor,
+    spatial_path: str,
+    temporal_path: str,
+    smooth_path: str,
+    device,
+    mode: str,
+    do_linear_blend: bool,
+):
+    spatial_net, temporal_net, smooth_net = load_models(
+        spatial_path, temporal_path, smooth_path, device
+    )
+
+    video_1 = video_preprocess(src_video_1[0].unsqueeze(0), img_h, img_w).to(device)
+    video_2 = video_preprocess(src_video_1, img_h, img_w).to(device)
+
+    video_list_2 = video_to_list_ts(video_2)
+
+    num_frames = len(video_list_2)
+    # motion estimation
+    smotion_tensor_list, omask_tensor_list = spatial_warping_single(
+        spatial_net, video_1, video_list_2, num_frames
+    )
+
+    # step 2: temporal warp
+    with torch.no_grad():
+        temporal_batch_out = build_TemporalNet(temporal_net, video_list_2)
+    tmotion_tensor_list = temporal_batch_out["motion_list"]
+
+    ##############################################
+    #############   data preparation  ############
+    # converting tmotion (t-th frame) into tsmotion ( (t-1)-th frame )
+
+    smesh_list, tsmotion_list = temporaling(smotion_tensor_list, tmotion_tensor_list)
+
+    # step 3: smooth warp
+    target_mesh = smooth_meshing(
+        smooth_net,
+        smotion_tensor_list,
+        omask_tensor_list,
+        tmotion_tensor_list,
+        smesh_list,
+        tsmotion_list,
+    )
+
+    stable_list, mask_list_2, img_2_list, _, _ = get_stable_sqe(
+        video_1, video_list_2, target_mesh, mode, do_linear_blend, is_single=True
+    )
+
+    stable_list = [stable.unsqueeze(0)[:, [2, 1, 0], :, :] for stable in stable_list]
+
+    return stable_list, img_2_list, mask_list_2
+
+
+def smooth_meshing(
+    smooth_net,
+    smotion_tensor_list,
+    omask_tensor_list,
+    tmotion_tensor_list,
+    smesh_list,
+    tsmotion_list,
+):
+    ori_mesh = 0
+    target_mesh = 0
+
+    for k in tqdm.tqdm(range(len(tmotion_tensor_list) - 6), "Smooth meshing"):
+        # get sublist and set the first element to 0
+        tsmotion_sublist = tsmotion_list[k : k + 7]
+        tsmotion_sublist[0] = smotion_tensor_list[k] * 0
+
+        with torch.no_grad():
+            smooth_batch_out = build_SmoothNet(
+                smooth_net,
+                tsmotion_sublist,
+                smesh_list[k : k + 7],
+                omask_tensor_list[k : k + 7],
+            )
+
+        _ori_mesh = smooth_batch_out["ori_mesh"]
+        _target_mesh = smooth_batch_out["target_mesh"]
+
+        if k == 0:
+            ori_mesh = _ori_mesh
+            target_mesh = _target_mesh
+
+        else:
+            ori_mesh = torch.cat((ori_mesh, _ori_mesh[:, -1, ...].unsqueeze(1)), 1)
+            target_mesh = torch.cat(
+                (target_mesh, _target_mesh[:, -1, ...].unsqueeze(1)), 1
+            )
+
+    return target_mesh
+
+
+def temporaling(smotion_tensor_list, tmotion_tensor_list):
+    smesh_list = []
+    tsmotion_list = []
+
+    rigid_mesh = get_rigid_mesh(1, img_h, img_w)
+    norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w)
+
+    for k in tqdm.tqdm(range(len(tmotion_tensor_list)), desc="Temporal calculating"):
+        smotion = smotion_tensor_list[k]
+        smesh = rigid_mesh + smotion
+        if k == 0:
+            tsmotion = smotion.clone() * 0
+
+        else:
+            smotion_1 = smotion_tensor_list[k - 1]
+            smesh_1 = rigid_mesh + smotion_1
+            tmotion = tmotion_tensor_list[k]
+            tmesh = rigid_mesh + tmotion
+            norm_smesh_1 = get_norm_mesh(smesh_1, img_h, img_w)
+            norm_tmesh = get_norm_mesh(tmesh, img_h, img_w)
+            tsmesh = torch_tps_transform_point.transformer(
+                norm_tmesh, norm_rigid_mesh, norm_smesh_1
+            )
+            tsmotion = recover_mesh(tsmesh, img_h, img_w) - smesh
+        # append
+        smesh_list.append(smesh)
+        tsmotion_list.append(tsmotion)
+    return smesh_list, tsmotion_list
+
+
+def spatial_warping(spatial_net, video_list_1, video_list_2, num_frames):
+    smotion_tensor_list = []
+    omask_tensor_list = []
+
+    for k in tqdm.tqdm(range(num_frames), desc="Motion estimating"):
+        # step 1: spatial warp
+        with torch.no_grad():
+            spatial_batch_out = build_SpatialNet(
+                spatial_net, video_list_1[k], video_list_2[k]
+            )
+        smotion = spatial_batch_out["motion"]
+        omask = spatial_batch_out["overlap_mesh"]
+        smotion_tensor_list.append(smotion)
+        omask_tensor_list.append(omask)
+    return smotion_tensor_list, omask_tensor_list
+
+
+def spatial_warping_single(spatial_net, video_1, video_list_2, num_frames):
+    smotion_tensor_list = []
+    omask_tensor_list = []
+
+    for k in tqdm.tqdm(range(num_frames), desc="Motion estimating"):
+        # step 1: spatial warp
+        with torch.no_grad():
+            spatial_batch_out = build_SpatialNet(spatial_net, video_1, video_list_2[k])
+        smotion = spatial_batch_out["motion"]
+        omask = spatial_batch_out["overlap_mesh"]
+        smotion_tensor_list.append(smotion)
+        omask_tensor_list.append(omask)
+    return smotion_tensor_list, omask_tensor_list
+
+
+def linear_blender(ref, tgt, ref_m, tgt_m, mask=False):
+    blur = GaussianBlur(kernel_size=(21, 21), sigma=20)
+    r1, c1 = torch.nonzero(ref_m[0, 0], as_tuple=True)
+    r2, c2 = torch.nonzero(tgt_m[0, 0], as_tuple=True)
+
+    center1 = (r1.float().mean(), c1.float().mean())
+    center2 = (r2.float().mean(), c2.float().mean())
+
+    vec = (center2[0] - center1[0], center2[1] - center1[1])
+
+    ovl = (ref_m * tgt_m).round()[:, 0].unsqueeze(1)
+    ref_m_ = ref_m[:, 0].unsqueeze(1) - ovl
+    r, c = torch.nonzero(ovl[0, 0], as_tuple=True)
+
+    ovl_mask = torch.zeros_like(ref_m_).cuda()
+    proj_val = (r - center1[0]) * vec[0] + (c - center1[1]) * vec[1]
+    ovl_mask[ovl.bool()] = (proj_val - proj_val.min()) / (
+        proj_val.max() - proj_val.min() + 1e-3
+    )
+
+    mask1 = (
+        blur(ref_m_ + (1 - ovl_mask) * ref_m[:, 0].unsqueeze(1)) * ref_m + ref_m_
+    ).clamp(0, 1)
+    if mask:
+        return mask1
+
+    mask2 = (1 - mask1) * tgt_m
+    stit = ref * mask1 + tgt * mask2
+
+    return stit
+
+
+def recover_mesh(norm_mesh, height, width):
+    # from [bs, pn, 2] to [bs, grid_h+1, grid_w+1, 2]
+
+    batch_size = norm_mesh.size()[0]
+    mesh_w = (norm_mesh[..., 0] + 1) * float(width) / 2.0
+    mesh_h = (norm_mesh[..., 1] + 1) * float(height) / 2.0
+    mesh = torch.stack([mesh_w, mesh_h], 2)  # [bs,(grid_h+1)*(grid_w+1),2]
+
+    return mesh.reshape([batch_size, grid_h + 1, grid_w + 1, 2])
+
+
+def get_rigid_mesh(batch_size, height, width):
+    ww = torch.matmul(
+        torch.ones([grid_h + 1, 1]),
+        torch.unsqueeze(torch.linspace(0.0, float(width), grid_w + 1), 0),
+    )
+    hh = torch.matmul(
+        torch.unsqueeze(torch.linspace(0.0, float(height), grid_h + 1), 1),
+        torch.ones([1, grid_w + 1]),
+    )
+    if torch.cuda.is_available():
+        ww = ww.cuda()
+        hh = hh.cuda()
+
+    ori_pt = torch.cat((ww.unsqueeze(2), hh.unsqueeze(2)), 2)  # (grid_h+1)*(grid_w+1)*2
+    ori_pt = ori_pt.unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+    return ori_pt
+
+
+def get_norm_mesh(mesh, height, width):
+    batch_size = mesh.size()[0]
+    mesh_w = mesh[..., 0] * 2.0 / float(width) - 1.0
+    mesh_h = mesh[..., 1] * 2.0 / float(height) - 1.0
+    norm_mesh = torch.stack([mesh_w, mesh_h], 3)  # bs*(grid_h+1)*(grid_w+1)*2
+
+    return norm_mesh.reshape([batch_size, -1, 2])  # bs*-1*2
+
+
+# bs, T, h, w, 2  smooth_path
+def get_stable_sqe(
+    img1_list, img2_list, ori_mesh, mode="FAST", do_linear_blend=False, is_single=False
+):
+    norm_rigid_mesh, width_min, height_min, out_width, out_height = get_new_sizes(
+        img2_list, ori_mesh
+    )
+
+    img1_warp = torch.zeros((1, 3, out_height, out_width)).cuda()
+
+    stable_list = []
+    mask_list_2 = []
+    img_2_list = []
+    for i in range(len(img2_list)):
+        if is_single:
+            img1 = (img1_list.cuda() + 1) / 2
+        else:
+            img1 = (img1_list[i].cuda() + 1) / 2
+        img2 = (img2_list[i].cuda() + 1) / 2
+        mesh = ori_mesh[:, i, :, :, :]
+        img2_warp, mask2, ave_fusion = get_fused(
+            mode,
+            do_linear_blend,
+            norm_rigid_mesh,
+            width_min,
+            height_min,
+            out_width,
+            out_height,
+            img1_warp,
+            img1,
+            img2,
+            mesh,
+        )
+
+        img_2_list.append(img2_warp)
+        mask_list_2.append(mask2)
+        stable_list.append(ave_fusion)
+
+    return stable_list, mask_list_2, img_2_list, out_width, out_height
+
+
+def get_new_sizes(img2_list, ori_mesh):
+    batch_size, _, img_h, img_w = img2_list[0].shape
+
+    rigid_mesh = get_rigid_mesh(batch_size, img_h, img_w)
+    norm_rigid_mesh = get_norm_mesh(rigid_mesh, img_h, img_w)
+
+    width_max = torch.max(torch.max(ori_mesh[..., 0]))
+    width_max = torch.maximum(torch.tensor(img_w).cuda(), width_max)
+    width_min = torch.min(ori_mesh[..., 0])
+    width_min = torch.minimum(torch.tensor(0).cuda(), width_min)
+    height_max = torch.max(ori_mesh[..., 1])
+    height_max = torch.maximum(torch.tensor(img_h).cuda(), height_max)
+    height_min = torch.min(ori_mesh[..., 1])
+    height_min = torch.minimum(torch.tensor(0).cuda(), height_min)
+
+    width_min = width_min.int()
+    width_max = width_max.int()
+    height_min = height_min.int()
+    height_max = height_max.int()
+    out_width = width_max - width_min + 1
+    out_height = height_max - height_min + 1
+    return norm_rigid_mesh, width_min, height_min, out_width, out_height
+
+
+def get_fused(
+    mode,
+    do_linear_blend,
+    norm_rigid_mesh,
+    width_min,
+    height_min,
+    out_width,
+    out_height,
+    img1_warp,
+    img1,
+    img2,
+    mesh,
+):
+    mesh_trans = torch.stack([mesh[..., 0] - width_min, mesh[..., 1] - height_min], 3)
+    norm_mesh = get_norm_mesh(mesh_trans, out_height, out_width)
+
+    # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary.
+    # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary.
+    img2_warp = torch_tps_transform.transformer(
+        img2, norm_mesh, norm_rigid_mesh, (out_height, out_width), mode=mode
+    )
+    mask2 = torch_tps_transform.transformer(
+        torch.ones_like(img2).cuda(),
+        norm_mesh,
+        norm_rigid_mesh,
+        (out_height, out_width),
+    )
+
+    if not do_linear_blend:
+        # average blending
+        img1_warp[
+            :,
+            :,
+            int(0 - height_min) : int(0 - height_min) + 360,
+            int(0 - width_min) : int(0 - width_min) + 480,
+        ] = img1
+
+        comp_1 = img1_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6)
+        comp_2 = img2_warp[0] / (img1_warp[0] + img2_warp[0] + 1e-6)
+
+        ave_fusion = img1_warp[0] * comp_1 + img2_warp[0] * comp_2
+    else:
+        img1_warp[
+            :,
+            :,
+            int(0 - height_min) : int(0 - height_min) + 360,
+            int(0 - width_min) : int(0 - width_min) + 480,
+        ] = img1
+        mask1 = img1_warp.clone()
+        mask1[
+            :,
+            :,
+            int(0 - height_min) : int(0 - height_min) + 360,
+            int(0 - width_min) : int(0 - width_min) + 480,
+        ] = 1
+        mask2 = torch_tps_transform.transformer(
+            torch.ones_like(img2).cuda(),
+            norm_mesh,
+            norm_rigid_mesh,
+            (out_height, out_width),
+        )
+        ave_fusion = linear_blender(img1_warp, img2_warp, mask1, mask2)
+        ave_fusion = ave_fusion[0]
+    return img2_warp, mask2, ave_fusion
+
+
+def load_models(spatial_path, temporal_path, smooth_path, device):
+    spatial_net = SpatialNet().to(device)
+    temporal_net = TemporalNet().to(device)
+    smooth_net = SmoothNet().to(device)
+
+    spatial_net.load_state_dict(torch.load(spatial_path)["model"])
+    temporal_net.load_state_dict(torch.load(temporal_path)["model"])
+    smooth_net.load_state_dict(torch.load(smooth_path)["model"])
+
+    spatial_net.eval()
+    temporal_net.eval()
+    smooth_net.eval()
+    return spatial_net, temporal_net, smooth_net

module_stabstitch/temporal_network.py

@@ -0,0 +1,410 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+from . import torch_DLT, torch_homo_transform
+
+grid_h = 6
+grid_w = 8
+
+
+# Covert global homo into mesh
+def H2Mesh(H, rigid_mesh):
+    H_inv = torch.inverse(H)
+    ori_pt = rigid_mesh.reshape(rigid_mesh.size()[0], -1, 2)
+    ones = torch.ones(rigid_mesh.size()[0], (grid_h + 1) * (grid_w + 1), 1)
+    if torch.cuda.is_available():
+        ori_pt = ori_pt.cuda()
+        ones = ones.cuda()
+
+    ori_pt = torch.cat((ori_pt, ones), 2)  # bs*(grid_h+1)*(grid_w+1)*3
+    tar_pt = torch.matmul(H_inv, ori_pt.permute(0, 2, 1))  # bs*3*(grid_h+1)*(grid_w+1)
+
+    mesh_x = torch.unsqueeze(tar_pt[:, 0, :] / tar_pt[:, 2, :], 2)
+    mesh_y = torch.unsqueeze(tar_pt[:, 1, :] / tar_pt[:, 2, :], 2)
+    mesh = torch.cat((mesh_x, mesh_y), 2).reshape(
+        [rigid_mesh.size()[0], grid_h + 1, grid_w + 1, 2]
+    )
+
+    return mesh
+
+
+def get_rigid_mesh(batch_size, height, width):
+    ww = torch.matmul(
+        torch.ones([grid_h + 1, 1]),
+        torch.unsqueeze(torch.linspace(0.0, float(width), grid_w + 1), 0),
+    )
+    hh = torch.matmul(
+        torch.unsqueeze(torch.linspace(0.0, float(height), grid_h + 1), 1),
+        torch.ones([1, grid_w + 1]),
+    )
+    if torch.cuda.is_available():
+        ww = ww.cuda()
+        hh = hh.cuda()
+
+    ori_pt = torch.cat((ww.unsqueeze(2), hh.unsqueeze(2)), 2)  # (grid_h+1)*(grid_w+1)*2
+    ori_pt = ori_pt.unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+    return ori_pt
+
+
+def get_norm_mesh(mesh, height, width):
+    batch_size = mesh.size()[0]
+    mesh_w = mesh[..., 0] * 2.0 / float(width) - 1.0
+    mesh_h = mesh[..., 1] * 2.0 / float(height) - 1.0
+    norm_mesh = torch.stack([mesh_w, mesh_h], 3)  # bs*(grid_h+1)*(grid_w+1)*2
+
+    return norm_mesh.reshape([batch_size, -1, 2])  # bs*-1*2
+
+
+def build_TemporalNet(net, img_tensor_list):
+    batch_size, _, img_h, img_w = img_tensor_list[0].size()
+    frame_num = len(img_tensor_list)
+
+    # for i in range(len(img_tensor_list)):
+    #     img_tensor_list[i] = (img_tensor_list[i] / 127.5) - 1.0
+
+    Homo_motion_list, Mesh_motion_list = net(img_tensor_list)
+
+    # output_homo_list = []
+    # output_mesh_list = []
+    # mesh_list = []
+    motion_list = []
+    for i in range(frame_num - 1):
+        # img1 = img_tensor_list[i]
+        # img2 = img_tensor_list[i+1]
+        Homo_motion = Homo_motion_list[i]
+        Mesh_motion = Mesh_motion_list[i]
+
+        # initialize the source points bs x 4 x 2
+        src_p = torch.tensor(
+            [[0.0, 0.0], [img_w - 1, 0.0], [0.0, img_h - 1], [img_w - 1, img_h - 1]]
+        )
+        if torch.cuda.is_available():
+            src_p = src_p.cuda()
+        src_p = src_p.unsqueeze(0).expand(batch_size, -1, -1)
+        # target points
+        dst_p = src_p + Homo_motion
+        # solve homo using DLT
+        H = torch_DLT.tensor_DLT(src_p, dst_p)
+
+        # mesh
+        rigid_mesh = get_rigid_mesh(batch_size, img_h, img_w)
+        ini_mesh = H2Mesh(H, rigid_mesh)
+        mesh = ini_mesh + Mesh_motion
+
+        motion_list.append(mesh - rigid_mesh)
+
+    motion_list.insert(0, torch.zeros([batch_size, grid_h + 1, grid_w + 1, 2]).cuda())
+    out_dict = {}
+    out_dict.update(motion_list=motion_list)
+
+    return out_dict
+
+
+def get_res18_FeatureMap(resnet18_model):
+    # feature extractor
+    # suppose input resolution to be 256*256   |   512*384
+    # the output resolution should be 16*16    |    32*24
+
+    layers_list = []
+
+    # mask
+    # layers_list.append(nn.Conv2d(6, 3, kernel_size=3, padding=1, bias=False))
+
+    layers_list.append(resnet18_model.conv1)  # stride 2*2     H/2
+    # layers_list.append(nn.Conv2d(6, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False))
+    layers_list.append(resnet18_model.bn1)
+    layers_list.append(resnet18_model.relu)
+    layers_list.append(resnet18_model.maxpool)  # stride 2       H/4
+
+    layers_list.append(resnet18_model.layer1)  # H/4
+    layers_list.append(resnet18_model.layer2)  # H/8
+
+    feature_extractor_stage1 = nn.Sequential(*layers_list)
+    feature_extractor_stage2 = nn.Sequential(resnet18_model.layer3)
+
+    return feature_extractor_stage1, feature_extractor_stage2
+
+
+# define and forward ( Because of the load is unbalanced when use torch.nn.DataParallel, we define warp in forward)
+class TemporalNet(nn.Module):
+    def __init__(self, dropout=0.0):
+        super(TemporalNet, self).__init__()
+
+        self.regressNet1_part1 = nn.Sequential(
+            nn.Conv2d(169, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 12, 20
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 6, 10
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 3, 5
+        )
+
+        self.regressNet1_part2 = nn.Sequential(
+            nn.Linear(in_features=768, out_features=512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=512, out_features=128, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=128, out_features=8, bias=True),
+        )
+
+        self.regressNet2_part1 = nn.Sequential(
+            nn.Conv2d(49, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 23, 40
+            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 12, 20
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 6, 10
+            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            # 3, 5
+        )
+
+        self.regressNet2_part2 = nn.Sequential(
+            nn.Linear(in_features=1536, out_features=1024, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_features=1024, out_features=512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(
+                in_features=512, out_features=(grid_w + 1) * (grid_h + 1) * 2, bias=True
+            ),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+        # ssl._create_default_https_context = ssl._create_unverified_context
+        resnet18_model = models.resnet.resnet18(weights="DEFAULT")
+        if torch.cuda.is_available():
+            resnet18_model = resnet18_model.cuda()
+        self.feature_extractor_stage1, self.feature_extractor_stage2 = (
+            get_res18_FeatureMap(resnet18_model)
+        )
+        # exit(0)
+
+    # forward ( Because of the load is unbalanced when use torch.nn.DataParallel, we define warp in forward)
+    def forward(self, img_tensor_list):
+        batch_size, _, img_h, img_w = img_tensor_list[0].size()
+        frame_num = len(img_tensor_list)
+
+        # extract feature maps for two resolutions
+
+        # feature1_list = []
+        # feature2_list = []
+        # for i in range(0, frame_num-1):
+        #     feature1_list.append(self.feature_extractor_stage1(img_tensor_list[i])) # bs, c, 45, 80
+        #     feature2_list.append(self.feature_extractor_stage2(feature1_list[i])) # bs, c, 23, 40
+
+        # calculate cost volume and regress motions
+        Homo_motion_list = []
+        Mesh_motion_list = []
+
+        feature1_list = [0, 0]
+        feature2_list = [0, 0]
+        for i in range(0, frame_num - 1):
+            if i == 0:
+                feature1_list[0] = self.feature_extractor_stage1(
+                    img_tensor_list[0].cuda()
+                )  # bs, c, 45, 80
+                feature2_list[0] = self.feature_extractor_stage2(feature1_list[0])
+
+                feature1_list[1] = self.feature_extractor_stage1(
+                    img_tensor_list[1].cuda()
+                )  # bs, c, 45, 80
+                feature2_list[1] = self.feature_extractor_stage2(feature1_list[1])
+            else:
+                feature1_list[1] = self.feature_extractor_stage1(
+                    img_tensor_list[i + 1].cuda()
+                )  # bs, c, 45, 80
+                feature2_list[1] = self.feature_extractor_stage2(feature1_list[1])
+
+            cv1 = self.cost_volume(
+                feature2_list[0], feature2_list[1], search_range=6, norm=False
+            )  # bs, 17^2, 23, 40
+            temp_1 = self.regressNet1_part1(cv1)
+            # print(temp_1.shape)
+            temp_1 = temp_1.view(temp_1.size()[0], -1)
+            offset_1 = self.regressNet1_part2(temp_1)
+            H_motion_1 = offset_1.reshape(-1, 4, 2)
+
+            src_p = torch.tensor(
+                [[0.0, 0.0], [img_w - 1, 0.0], [0.0, img_h - 1], [img_w - 1, img_h - 1]]
+            )
+            if torch.cuda.is_available():
+                src_p = src_p.cuda()
+            src_p = src_p.unsqueeze(0).expand(batch_size, -1, -1)
+            dst_p = src_p + H_motion_1
+            H = torch_DLT.tensor_DLT(src_p / 8, dst_p / 8)
+
+            M_tensor = torch.tensor(
+                [
+                    [img_w / 8 / 2.0, 0.0, img_w / 8 / 2.0],
+                    [0.0, img_h / 8 / 2.0, img_h / 8 / 2.0],
+                    [0.0, 0.0, 1.0],
+                ]
+            )
+
+            if torch.cuda.is_available():
+                M_tensor = M_tensor.cuda()
+
+            M_tile = M_tensor.unsqueeze(0).expand(batch_size, -1, -1)
+            M_tensor_inv = torch.inverse(M_tensor)
+            M_tile_inv = M_tensor_inv.unsqueeze(0).expand(batch_size, -1, -1)
+            H_mat = torch.matmul(torch.matmul(M_tile_inv, H), M_tile)
+
+            warp_feature2 = torch_homo_transform.transformer(
+                feature1_list[1], H_mat, (int(img_h / 8), int(img_w / 8))
+            )
+
+            ######### stage 2
+            cv2 = self.cost_volume(
+                feature1_list[0], warp_feature2, search_range=3, norm=False
+            )  # bs, 9^2, 45, 80
+            temp_2 = self.regressNet2_part1(cv2)
+            temp_2 = temp_2.view(temp_2.size()[0], -1)
+            offset_2 = self.regressNet2_part2(temp_2)
+            M_motion_2 = offset_2.reshape(-1, grid_h + 1, grid_w + 1, 2)
+
+            Homo_motion_list.append(H_motion_1)
+            Mesh_motion_list.append(M_motion_2)
+            feature1_list[0] = feature1_list[1]
+            feature2_list[0] = feature2_list[1]
+
+        return Homo_motion_list, Mesh_motion_list
+
+    @staticmethod
+    def cost_volume(x1, x2, search_range, norm=True, fast=True):
+        if norm:
+            x1 = F.normalize(x1, p=2, dim=1)
+            x2 = F.normalize(x2, p=2, dim=1)
+        bs, c, h, w = x1.shape
+        padded_x2 = F.pad(x2, [search_range] * 4)  # [b,c,h,w] -> [b,c,h+sr*2,w+sr*2]
+        max_offset = search_range * 2 + 1
+
+        if fast:
+            # faster(*2) but cost higher(*n) GPU memory
+            patches = F.unfold(padded_x2, (max_offset, max_offset)).reshape(
+                bs, c, max_offset**2, h, w
+            )
+            cost_vol = (x1.unsqueeze(2) * patches).mean(dim=1, keepdim=False)
+        else:
+            # slower but save memory
+            cost_vol = []
+            for j in range(0, max_offset):
+                for i in range(0, max_offset):
+                    x2_slice = padded_x2[:, :, j : j + h, i : i + w]
+                    cost = torch.mean(x1 * x2_slice, dim=1, keepdim=True)
+                    cost_vol.append(cost)
+            cost_vol = torch.cat(cost_vol, dim=1)
+
+        cost_vol = F.leaky_relu(cost_vol, 0.1)
+
+        return cost_vol
+
+    def extract_patches(self, x, kernel=3, stride=1):
+        if kernel != 1:
+            x = nn.ZeroPad2d(1)(x)
+        x = x.permute(0, 2, 3, 1)
+        all_patches = x.unfold(1, kernel, stride).unfold(2, kernel, stride)
+        return all_patches
+
+    def CCL(self, feature_1, feature_2):
+        bs, c, h, w = feature_1.size()
+
+        norm_feature_1 = F.normalize(feature_1, p=2, dim=1)
+        norm_feature_2 = F.normalize(feature_2, p=2, dim=1)
+        # print(norm_feature_2.size())
+
+        patches = self.extract_patches(norm_feature_2)
+        if torch.cuda.is_available():
+            patches = patches.cuda()
+
+        matching_filters = patches.reshape(
+            (
+                patches.size()[0],
+                -1,
+                patches.size()[3],
+                patches.size()[4],
+                patches.size()[5],
+            )
+        )
+
+        match_vol = []
+        for i in range(bs):
+            single_match = F.conv2d(
+                norm_feature_1[i].unsqueeze(0), matching_filters[i], padding=1
+            )
+            match_vol.append(single_match)
+
+        match_vol = torch.cat(match_vol, 0)
+        # print(match_vol .size())
+
+        # scale softmax
+        softmax_scale = 10
+        match_vol = F.softmax(match_vol * softmax_scale, 1)
+
+        channel = match_vol.size()[1]
+
+        h_one = torch.linspace(0, h - 1, h)
+        one1w = torch.ones(1, w)
+        if torch.cuda.is_available():
+            h_one = h_one.cuda()
+            one1w = one1w.cuda()
+        h_one = torch.matmul(h_one.unsqueeze(1), one1w)
+        h_one = h_one.unsqueeze(0).unsqueeze(0).expand(bs, channel, -1, -1)
+
+        w_one = torch.linspace(0, w - 1, w)
+        oneh1 = torch.ones(h, 1)
+        if torch.cuda.is_available():
+            w_one = w_one.cuda()
+            oneh1 = oneh1.cuda()
+        w_one = torch.matmul(oneh1, w_one.unsqueeze(0))
+        w_one = w_one.unsqueeze(0).unsqueeze(0).expand(bs, channel, -1, -1)
+
+        c_one = torch.linspace(0, channel - 1, channel)
+        if torch.cuda.is_available():
+            c_one = c_one.cuda()
+        c_one = c_one.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand(bs, -1, h, w)
+
+        flow_h = match_vol * (c_one // w - h_one)
+        flow_h = torch.sum(flow_h, dim=1, keepdim=True)
+        flow_w = match_vol * (c_one % w - w_one)
+        flow_w = torch.sum(flow_w, dim=1, keepdim=True)
+
+        feature_flow = torch.cat([flow_w, flow_h], 1)
+        # print(flow.size())
+
+        return feature_flow

module_stabstitch/torch_DLT.py

@@ -0,0 +1,43 @@
+import torch
+
+# src_p: shape=(bs, 4, 2)
+# det_p: shape=(bs, 4, 2)
+#
+#                                     | h1 |
+#                                     | h2 |
+#                                     | h3 |
+# | x1 y1 1  0  0  0  -x1x2  -y1x2 |  | h4 |  =  | x2 |
+# | 0  0  0  x1 y1 1  -x1y2  -y1y2 |  | h5 |     | y2 |
+#                                     | h6 |
+#                                     | h7 |
+#                                     | h8 |
+
+
+def tensor_DLT(src_p, dst_p):
+    bs, _, _ = src_p.shape
+
+    ones = torch.ones(bs, 4, 1)
+    if torch.cuda.is_available():
+        ones = ones.cuda()
+    xy1 = torch.cat((src_p, ones), 2)
+    zeros = torch.zeros_like(xy1)
+    if torch.cuda.is_available():
+        zeros = zeros.cuda()
+
+    xyu, xyd = torch.cat((xy1, zeros), 2), torch.cat((zeros, xy1), 2)
+    M1 = torch.cat((xyu, xyd), 2).reshape(bs, -1, 6)
+    M2 = torch.matmul(
+        dst_p.reshape(-1, 2, 1),
+        src_p.reshape(-1, 1, 2),
+    ).reshape(bs, -1, 2)
+
+    # Ah = b
+    A = torch.cat((M1, -M2), 2)
+    b = dst_p.reshape(bs, -1, 1)
+
+    # h = A^{-1}b
+    Ainv = torch.inverse(A)
+    h8 = torch.matmul(Ainv, b).reshape(bs, 8)
+
+    H = torch.cat((h8, ones[:, 0, :]), 1).reshape(bs, 3, 3)
+    return H

module_stabstitch/torch_homo_transform.py

@@ -0,0 +1,189 @@
+import numpy as np
+import torch
+
+
+def transformer(U, theta, out_size, **kwargs):
+    """Spatial Transformer Layer
+
+    Implements a spatial transformer layer as described in [1]_.
+    Based on [2]_ and edited by David Dao for Tensorflow.
+
+    Parameters
+    ----------
+    U : float
+        The output of a convolutional net should have the
+        shape [num_batch, height, width, num_channels].
+    theta: float
+        The output of the
+        localisation network should be [num_batch, 6].
+    out_size: tuple of two ints
+        The size of the output of the network (height, width)
+
+    References
+    ----------
+    .. [1]  Spatial Transformer Networks
+            Max Jaderberg, Karen Simonyan, Andrew Zisserman, Koray Kavukcuoglu
+            Submitted on 5 Jun 2015
+    .. [2]  https://github.com/skaae/transformer_network/blob/master/transformerlayer.py
+
+    Notes
+    -----
+    To initialize the network to the identity transform init
+    ``theta`` to :
+        identity = np.array([[1., 0., 0.],
+                             [0., 1., 0.]])
+        identity = identity.flatten()
+        theta = tf.Variable(initial_value=identity)
+
+    """
+
+    def _repeat(x, n_repeats):
+        rep = torch.ones(
+            [
+                n_repeats,
+            ]
+        ).unsqueeze(0)
+        rep = rep.int()
+        x = x.int()
+
+        x = torch.matmul(x.reshape([-1, 1]), rep)
+        return x.reshape([-1])
+
+    def _interpolate(im, x, y, out_size):
+        num_batch, num_channels, height, width = im.size()
+
+        height_f = height
+        width_f = width
+        out_height, out_width = out_size[0], out_size[1]
+
+        zero = 0
+        max_y = height - 1
+        max_x = width - 1
+
+        x = (x + 1.0) * (width_f) / 2.0
+        y = (y + 1.0) * (height_f) / 2.0
+
+        # do sampling
+        x0 = torch.floor(x).int()
+        x1 = x0 + 1
+        y0 = torch.floor(y).int()
+        y1 = y0 + 1
+
+        x0 = torch.clamp(x0, zero, max_x)
+        x1 = torch.clamp(x1, zero, max_x)
+        y0 = torch.clamp(y0, zero, max_y)
+        y1 = torch.clamp(y1, zero, max_y)
+        dim2 = torch.from_numpy(np.array(width))
+        dim1 = torch.from_numpy(np.array(width * height))
+
+        base = _repeat(torch.arange(0, num_batch) * dim1, out_height * out_width)
+        if torch.cuda.is_available():
+            dim2 = dim2.cuda()
+            dim1 = dim1.cuda()
+            y0 = y0.cuda()
+            y1 = y1.cuda()
+            x0 = x0.cuda()
+            x1 = x1.cuda()
+            base = base.cuda()
+        base_y0 = base + y0 * dim2
+        base_y1 = base + y1 * dim2
+        idx_a = base_y0 + x0
+        idx_b = base_y1 + x0
+        idx_c = base_y0 + x1
+        idx_d = base_y1 + x1
+
+        # channels dim
+        im = im.permute(0, 2, 3, 1)
+        im_flat = im.reshape([-1, num_channels]).float()
+
+        idx_a = idx_a.unsqueeze(-1).long()
+        idx_a = idx_a.expand(out_height * out_width * num_batch, num_channels)
+        Ia = torch.gather(im_flat, 0, idx_a)
+
+        idx_b = idx_b.unsqueeze(-1).long()
+        idx_b = idx_b.expand(out_height * out_width * num_batch, num_channels)
+        Ib = torch.gather(im_flat, 0, idx_b)
+
+        idx_c = idx_c.unsqueeze(-1).long()
+        idx_c = idx_c.expand(out_height * out_width * num_batch, num_channels)
+        Ic = torch.gather(im_flat, 0, idx_c)
+
+        idx_d = idx_d.unsqueeze(-1).long()
+        idx_d = idx_d.expand(out_height * out_width * num_batch, num_channels)
+        Id = torch.gather(im_flat, 0, idx_d)
+
+        x0_f = x0.float()
+        x1_f = x1.float()
+        y0_f = y0.float()
+        y1_f = y1.float()
+
+        wa = torch.unsqueeze(((x1_f - x) * (y1_f - y)), 1)
+        wb = torch.unsqueeze(((x1_f - x) * (y - y0_f)), 1)
+        wc = torch.unsqueeze(((x - x0_f) * (y1_f - y)), 1)
+        wd = torch.unsqueeze(((x - x0_f) * (y - y0_f)), 1)
+        output = wa * Ia + wb * Ib + wc * Ic + wd * Id
+
+        return output
+
+    def _meshgrid(height, width):
+        x_t = torch.matmul(
+            torch.ones([height, 1]),
+            torch.transpose(torch.unsqueeze(torch.linspace(-1.0, 1.0, width), 1), 1, 0),
+        )
+        y_t = torch.matmul(
+            torch.unsqueeze(torch.linspace(-1.0, 1.0, height), 1),
+            torch.ones([1, width]),
+        )
+        # x_t = torch.matmul(torch.ones([height, 1]),
+        #                       torch.transpose(torch.unsqueeze(torch.linspace(0.0, width.float(), width), 1), 1, 0))
+        # y_t = torch.matmul(torch.unsqueeze(torch.linspace(0.0, height.float(), height), 1),
+        #                       torch.ones([1, width]))
+
+        x_t_flat = x_t.reshape((1, -1)).float()
+        y_t_flat = y_t.reshape((1, -1)).float()
+
+        ones = torch.ones_like(x_t_flat)
+        grid = torch.cat([x_t_flat, y_t_flat, ones], 0)
+        if torch.cuda.is_available():
+            grid = grid.cuda()
+        return grid
+
+    def _transform(theta, input_dim, out_size):
+        num_batch, num_channels, height, width = input_dim.size()
+        #  Changed
+        theta = theta.reshape([-1, 3, 3]).float()
+
+        out_height, out_width = out_size[0], out_size[1]
+        grid = _meshgrid(out_height, out_width)
+        grid = grid.unsqueeze(0).reshape([1, -1])
+        shape = grid.size()
+        grid = grid.expand(num_batch, shape[1])
+        grid = grid.reshape([num_batch, 3, -1])
+
+        T_g = torch.matmul(theta, grid)
+        x_s = T_g[:, 0, :]
+        y_s = T_g[:, 1, :]
+        t_s = T_g[:, 2, :]
+
+        t_s_flat = t_s.reshape([-1])
+
+        # smaller
+        small = 1e-7
+        smallers = 1e-6 * (1.0 - torch.ge(torch.abs(t_s_flat), small).float())
+
+        t_s_flat = t_s_flat + smallers
+        # condition = torch.sum(torch.gt(torch.abs(t_s_flat), small).float())
+        # Ty changed
+        x_s_flat = x_s.reshape([-1]) / t_s_flat
+        y_s_flat = y_s.reshape([-1]) / t_s_flat
+
+        input_transformed = _interpolate(input_dim, x_s_flat, y_s_flat, out_size)
+
+        output = input_transformed.reshape(
+            [num_batch, out_height, out_width, num_channels]
+        )
+        output = output.permute(0, 3, 1, 2)
+        return output  # , condition
+
+    output = _transform(theta, U, out_size)
+    return output  # , condition

module_stabstitch/torch_tps_transform.py

@@ -0,0 +1,247 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def transformer(U, source, target, out_size, mode="NORMAL"):
+    """
+      Thin Plate Spline Spatial Transformer Layer
+    TPS control points are arranged in arbitrary positions given by `source`.
+    U : float Tensor [num_batch, height, width, num_channels].
+      Input Tensor.
+    source : float Tensor [num_batch, num_point, 2]
+      The source position of the control points.
+    target : float Tensor [num_batch, num_point, 2]
+      The target position of the control points.
+    out_size: tuple of two integers [height, width]
+      The size of the output of the network (height, width)
+    """
+
+    def _repeat(x, n_repeats):
+        rep = torch.ones(
+            [
+                n_repeats,
+            ]
+        ).unsqueeze(0)
+        rep = rep.int()
+        x = x.int()
+
+        x = torch.matmul(x.reshape([-1, 1]), rep)
+        return x.reshape([-1])
+
+    def _interpolate(im, x, y, out_size):
+        num_batch, num_channels, height, width = im.size()
+
+        height_f = height
+        width_f = width
+        out_height, out_width = out_size[0], out_size[1]
+
+        zero = 0
+        max_y = height - 1
+        max_x = width - 1
+
+        x = (x + 1.0) * (width_f) / 2.0
+        y = (y + 1.0) * (height_f) / 2.0
+
+        # do sampling
+        x0 = torch.floor(x).int()
+        x1 = x0 + 1
+        y0 = torch.floor(y).int()
+        y1 = y0 + 1
+
+        x0 = torch.clamp(x0, zero, max_x)
+        x1 = torch.clamp(x1, zero, max_x)
+        y0 = torch.clamp(y0, zero, max_y)
+        y1 = torch.clamp(y1, zero, max_y)
+        dim2 = torch.from_numpy(np.array(width))
+        dim1 = torch.from_numpy(np.array(width * height))
+
+        base = _repeat(torch.arange(0, num_batch) * dim1, out_height * out_width)
+        if torch.cuda.is_available():
+            dim2 = dim2.cuda()
+            dim1 = dim1.cuda()
+            y0 = y0.cuda()
+            y1 = y1.cuda()
+            x0 = x0.cuda()
+            x1 = x1.cuda()
+            base = base.cuda()
+        base_y0 = base + y0 * dim2
+        base_y1 = base + y1 * dim2
+        idx_a = base_y0 + x0
+        idx_b = base_y1 + x0
+        idx_c = base_y0 + x1
+        idx_d = base_y1 + x1
+
+        # channels dim
+        im = im.permute(0, 2, 3, 1)
+        im_flat = im.reshape([-1, num_channels]).float()
+
+        idx_a = idx_a.unsqueeze(-1).long()
+        idx_a = idx_a.expand(out_height * out_width * num_batch, num_channels)
+        Ia = torch.gather(im_flat, 0, idx_a)
+
+        idx_b = idx_b.unsqueeze(-1).long()
+        idx_b = idx_b.expand(out_height * out_width * num_batch, num_channels)
+        Ib = torch.gather(im_flat, 0, idx_b)
+
+        idx_c = idx_c.unsqueeze(-1).long()
+        idx_c = idx_c.expand(out_height * out_width * num_batch, num_channels)
+        Ic = torch.gather(im_flat, 0, idx_c)
+
+        idx_d = idx_d.unsqueeze(-1).long()
+        idx_d = idx_d.expand(out_height * out_width * num_batch, num_channels)
+        Id = torch.gather(im_flat, 0, idx_d)
+
+        x0_f = x0.float()
+        x1_f = x1.float()
+        y0_f = y0.float()
+        y1_f = y1.float()
+
+        wa = torch.unsqueeze(((x1_f - x) * (y1_f - y)), 1)
+        wb = torch.unsqueeze(((x1_f - x) * (y - y0_f)), 1)
+        wc = torch.unsqueeze(((x - x0_f) * (y1_f - y)), 1)
+        wd = torch.unsqueeze(((x - x0_f) * (y - y0_f)), 1)
+        output = wa * Ia + wb * Ib + wc * Ic + wd * Id
+
+        return output
+
+    def _meshgrid(height, width, source):
+        x_t = torch.matmul(
+            torch.ones([height, 1]),
+            torch.unsqueeze(torch.linspace(-1.0, 1.0, width), 0),
+        )
+        y_t = torch.matmul(
+            torch.unsqueeze(torch.linspace(-1.0, 1.0, height), 1),
+            torch.ones([1, width]),
+        )
+        if torch.cuda.is_available():
+            x_t = x_t.cuda()
+            y_t = y_t.cuda()
+
+        x_t_flat = x_t.reshape([1, 1, -1])
+        y_t_flat = y_t.reshape([1, 1, -1])
+
+        num_batch = source.size()[0]
+        px = torch.unsqueeze(source[:, :, 0], 2)  # [bn, pn, 1]
+        py = torch.unsqueeze(source[:, :, 1], 2)  # [bn, pn, 1]
+        if torch.cuda.is_available():
+            px = px.cuda()
+            py = py.cuda()
+        d2 = torch.square(x_t_flat - px) + torch.square(y_t_flat - py)
+        r = d2 * torch.log(d2 + 1e-6)  # [bn, pn, h*w]
+        x_t_flat_g = x_t_flat.expand(num_batch, -1, -1)  # [bn, 1, h*w]
+        y_t_flat_g = y_t_flat.expand(num_batch, -1, -1)  # [bn, 1, h*w]
+        ones = torch.ones_like(x_t_flat_g)  # [bn, 1, h*w]
+        if torch.cuda.is_available():
+            ones = ones.cuda()
+
+        grid = torch.cat((ones, x_t_flat_g, y_t_flat_g, r), 1)  # [bn, 3+pn, h*w]
+
+        # if torch.cuda.is_available():
+        #    grid = grid.cuda()
+        return grid
+
+    def _transform(T, source, input_dim, out_size, mode):
+        num_batch, num_channels, height, width = input_dim.size()
+
+        out_height, out_width = out_size[0], out_size[1]
+        grid = _meshgrid(out_height, out_width, source)  # [bn, 3+pn, h*w]
+
+        # transform A x (1, x_t, y_t, r1, r2, ..., rn) -> (x_s, y_s)
+        # [bn, 2, pn+3] x [bn, pn+3, h*w] -> [bn, 2, h*w]
+        T_g = torch.matmul(T, grid)
+        x_s = T_g[:, 0, :]
+        y_s = T_g[:, 1, :]
+
+        if mode == "NORMAL":
+            # original iterpolation implementation
+            x_s_flat = x_s.reshape([-1])
+            y_s_flat = y_s.reshape([-1])
+            input_transformed = _interpolate(input_dim, x_s_flat, y_s_flat, out_size)
+            output = input_transformed.reshape(
+                [num_batch, out_height, out_width, num_channels]
+            )
+            output = output.permute(0, 3, 1, 2)
+        else:
+            # using F.grid_sample to implement
+            x_s_flat = x_s.reshape([num_batch, 1, out_height, out_width])
+            y_s_flat = y_s.reshape([num_batch, 1, out_height, out_width])
+            output = F.grid_sample(
+                input_dim,
+                torch.cat([x_s_flat, y_s_flat], 1).permute(0, 2, 3, 1),
+                align_corners=True,
+            )
+
+        return output
+
+    def _solve_system(source, target):
+        num_batch = source.size()[0]
+        num_point = source.size()[1]
+
+        np.set_printoptions(precision=8)
+
+        ones = torch.ones(num_batch, num_point, 1).float()
+        if torch.cuda.is_available():
+            ones = ones.cuda()
+        p = torch.cat([ones, source], 2)  # [bn, pn, 3]
+
+        p_1 = p.reshape([num_batch, -1, 1, 3])  # [bn, pn, 1, 3]
+        p_2 = p.reshape([num_batch, 1, -1, 3])  # [bn, 1, pn, 3]
+        d2 = torch.sum(
+            torch.square(p_1 - p_2), 3
+        )  # p1 - p2: [bn, pn, pn, 3]   final output: [bn, pn, pn]
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # torch.set_printoptions(precision=8)
+        # print(d2[0])
+        # print(d2.dtype)
+        r = d2 * torch.log(d2 + 1e-6)  # [bn, pn, pn]
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print(torch.log(d2 + 1e-6)[0].cpu().detach().numpy())
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print(r[0].cpu().detach().numpy())
+        # print("---------------------------------------------------------")
+
+        zeros = torch.zeros(num_batch, 3, 3).float()
+        if torch.cuda.is_available():
+            zeros = zeros.cuda()
+        W_0 = torch.cat((p, r), 2)  # [bn, pn, 3+pn]
+        W_1 = torch.cat((zeros, p.permute(0, 2, 1)), 2)  # [bn, 3, pn+3]
+        W = torch.cat((W_0, W_1), 1)  # [bn, pn+3, pn+3]
+
+        # w = np.load("www.npy")
+        # w = torch.tensor(w).cuda()
+        # W = w.expand(source.size()[0],-1,-1)
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print("W")
+        # print(W[0].cpu().detach().numpy())
+        W_inv = torch.inverse(W.type(torch.float64))
+
+        # W_inv_np = np.linalg.inv(W.cpu().detach().numpy())
+        # W_inv = torch.tensor(W_inv_np).cuda()
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print("W_inv")
+        # print(W_inv[0].cpu().detach().numpy())
+
+        zeros2 = torch.zeros(num_batch, 3, 2)
+        if torch.cuda.is_available():
+            zeros2 = zeros2.cuda()
+        tp = torch.cat((target, zeros2), 1)  # [bn, pn+3, 2]
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print(tp[0].cpu().detach().numpy())
+        T = torch.matmul(W_inv, tp.type(torch.float64))  # [bn, pn+3, 2]
+        T = T.permute(0, 2, 1)  # [bn, 2, pn+3]
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # print(T[0].cpu().detach().numpy())
+        # print(T.size())
+
+        return T.type(torch.float32)
+
+    T = _solve_system(source, target)
+    # t = np.load("ttt.npy")
+    # t = torch.tensor(t).cuda()
+    # T = t.expand(source.size()[0],-1,-1)
+    # print("xxxxxxxxxxxxxxxxxxxx")
+    # print(T[0].cpu().detach().numpy())
+    output = _transform(T, source, U, out_size, mode)
+
+    return output  # , condition

module_stabstitch/torch_tps_transform_point.py

@@ -0,0 +1,123 @@
+import numpy as np
+import torch
+
+
+def transformer(point, source, target):
+    """
+      Thin Plate Spline Spatial Transformer Layer
+    TPS control points are arranged in arbitrary positions given by `source`.
+    U : float Tensor [num_batch, height, width, num_channels].
+      Input Tensor.
+    source : float Tensor [num_batch, num_point, 2]
+      The source position of the control points.
+    target : float Tensor [num_batch, num_point, 2]
+      The target position of the control points.
+    out_size: tuple of two integers [height, width]
+      The size of the output of the network (height, width)
+    """
+
+    # point: bn, num_point, 2
+    def _meshgrid_point(point, source):
+        # x_t = torch.matmul(torch.ones([height, 1]), torch.unsqueeze(torch.linspace(-1.0, 1.0, width), 0))
+        # y_t = torch.matmul(torch.unsqueeze(torch.linspace(-1.0, 1.0, height), 1), torch.ones([1, width]))
+        # if torch.cuda.is_available():
+        #     x_t = x_t.cuda()
+        #     y_t = y_t.cuda()
+
+        # x_t_flat = x_t.reshape([1, 1, -1])
+        # y_t_flat = y_t.reshape([1, 1, -1])
+
+        num_batch, num_point, _ = point.size()
+
+        x_t_flat = point[:, :, 0].view(num_batch, 1, num_point)  # bs, 1, num_point
+        y_t_flat = point[:, :, 1].view(num_batch, 1, num_point)
+
+        num_batch = source.size()[0]
+        px = torch.unsqueeze(source[:, :, 0], 2)  # [bn, pn, 1]
+        py = torch.unsqueeze(source[:, :, 1], 2)  # [bn, pn, 1]
+        if torch.cuda.is_available():
+            px = px.cuda()
+            py = py.cuda()
+        d2 = torch.square(x_t_flat - px) + torch.square(y_t_flat - py)
+        r = d2 * torch.log(d2 + 1e-6)  # [bn, pn, h*w]
+        # x_t_flat_g = x_t_flat.expand(num_batch, -1, -1)  # [bn, 1, h*w]
+        # y_t_flat_g = y_t_flat.expand(num_batch, -1, -1)  # [bn, 1, h*w]
+        ones = torch.ones_like(x_t_flat)  # [bn, 1, h*w]
+        if torch.cuda.is_available():
+            ones = ones.cuda()
+
+        grid = torch.cat((ones, x_t_flat, y_t_flat, r), 1)  # [bn, 3+pn, num_point]
+
+        return grid
+
+    def _transform(T, source, point):
+        num_batch, num_point, num_channels = point.size()
+
+        # out_height, out_width = out_size[0], out_size[1]
+        # grid = _meshgrid(out_height, out_width, source) # [bn, 3+pn, h*w]
+        grid = _meshgrid_point(point, source)  # [bn, 3+pn, num_point]
+
+        # transform A x (1, x_t, y_t, r1, r2, ..., rn) -> (x_s, y_s)
+        # [bn, 2, pn+3] x [bn, pn+3, h*w] -> [bn, 2, h*w]
+        T_g = torch.matmul(T, grid)  # [bn, 2, num_point]
+        # x_s = T_g[:,0,:]
+        # y_s = T_g[:,1,:]
+        # x_s_flat = x_s.reshape([-1])
+        # y_s_flat = y_s.reshape([-1])
+
+        # input_transformed = _interpolate(input_dim, x_s_flat, y_s_flat,out_size)
+
+        # #output = input_transformed.reshape([num_batch, out_height, out_width, num_channels])
+
+        # output = output.permute(0,3,1,2)
+        T_g = T_g.permute(0, 2, 1)
+        output = T_g.view(num_batch, num_point, num_channels)
+
+        return output  # , condition
+
+    def _solve_system(source, target):
+        num_batch = source.size()[0]
+        num_point = source.size()[1]
+
+        np.set_printoptions(precision=8)
+
+        ones = torch.ones(num_batch, num_point, 1).float()
+        if torch.cuda.is_available():
+            ones = ones.cuda()
+        p = torch.cat([ones, source], 2)  # [bn, pn, 3]
+
+        p_1 = p.reshape([num_batch, -1, 1, 3])  # [bn, pn, 1, 3]
+        p_2 = p.reshape([num_batch, 1, -1, 3])  # [bn, 1, pn, 3]
+        d2 = torch.sum(
+            torch.square(p_1 - p_2), 3
+        )  # p1 - p2: [bn, pn, pn, 3]   final output: [bn, pn, pn]
+        # print("xxxxxxxxxxxxxxxxxxxx")
+        # torch.set_printoptions(precision=8)
+        # print(d2[0])
+        # print(d2.dtype)
+        r = d2 * torch.log(d2 + 1e-6)  # [bn, pn, pn]
+
+        zeros = torch.zeros(num_batch, 3, 3).float()
+        if torch.cuda.is_available():
+            zeros = zeros.cuda()
+        W_0 = torch.cat((p, r), 2)  # [bn, pn, 3+pn]
+        W_1 = torch.cat((zeros, p.permute(0, 2, 1)), 2)  # [bn, 3, pn+3]
+        W = torch.cat((W_0, W_1), 1)  # [bn, pn+3, pn+3]
+
+        W_inv = torch.inverse(W.type(torch.float64))
+
+        zeros2 = torch.zeros(num_batch, 3, 2)
+        if torch.cuda.is_available():
+            zeros2 = zeros2.cuda()
+        tp = torch.cat((target, zeros2), 1)  # [bn, pn+3, 2]
+
+        T = torch.matmul(W_inv, tp.type(torch.float64))  # [bn, pn+3, 2]
+        T = T.permute(0, 2, 1)  # [bn, 2, pn+3]
+
+        return T.type(torch.float32)
+
+    T = _solve_system(source, target)
+
+    output = _transform(T, source, point)
+
+    return output  # , condition

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "fm_nodes"
-description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba, ConvIR"
-version = "1.0.5"
+description = "An expanding collection of ComfyUI nodes focused on new papers. Including: WFEN, RealViFormer, ProPIH, CoLIE, VFIMamba, ConvIR, StabStitch"
+version = "1.0.6"
 license = {file = "LICENSE"}
 dependencies = ["torch", "einops"]
 

run.py

@@ -2,7 +2,7 @@
 @author: Fuou Marinas
 @title: FM Nodes
 @nickname: FM_nodes
-@description: A collection of nodes. WFEN Face Super Resolution.
+@description: A collection of nodes.
 """
 
 from pathlib import Path
@@ -12,17 +12,17 @@ import torch.nn.functional as F
 import tqdm
 from comfy.utils import ProgressBar
 
-from .module_colie.colie_model import CoLIE_Config, run_colie
 from .constants import (
     PROPIH_G_MODEL_PATH,
     PROPIH_VGG_MODEL_PATH,
     REALVIFORMER_MODEL_PATH,
     WFEN_MODEL_PATH,
 )
+from .module_colie.colie_model import CoLIE_Config, run_colie
 from .module_propih.propih_model import VGG19HRNetModel
 from .module_realviformer.realviformer_arch import RealViformer
-from .utils import ensure_size, img_to_mask
 from .module_wfen.wfen_model import WFENModel
+from .utils import ensure_size, img_to_mask
 
 base_dir = Path(__file__).resolve().parent
 

run_stabstitch.py

@@ -0,0 +1,183 @@
+from pathlib import Path
+
+import torch
+
+from .constants import (
+    STAB_CONSTRUCT_DEFAULT,
+    STAB_CONSTRUCT_MODES,
+    STAB_SMOOTH_PATH,
+    STAB_SPATIAL_PATH,
+    STAB_TEMPORAL_PATH,
+)
+from .module_stabstitch.stabstitch_model import (
+    crop_or_resize,
+    stabstitch_inference,
+    stabstitch_single_inference,
+)
+
+base_dir = Path(__file__).resolve().parent
+
+
+class StabStitch_Stitch:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "video_1": ("IMAGE",),
+                "video_2": ("IMAGE",),
+                # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary.
+                # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary.
+                "interpolate_mode": (
+                    STAB_CONSTRUCT_MODES,
+                    {"default": STAB_CONSTRUCT_DEFAULT},
+                ),
+                "do_linear_blend": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+            },
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "IMAGE",
+        "IMAGE",
+    )
+    RETURN_NAMES = (
+        "combined",
+        "stab",
+        "mask",
+    )
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/StabStitch"
+
+    def todo(
+        self,
+        video_1: torch.Tensor,
+        video_2: torch.Tensor,
+        interpolate_mode: str,
+        do_linear_blend: bool,
+    ):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        stable_list, video_frames_2, mask_frames_2 = stabstitch_inference(
+            video_1,
+            video_2,
+            spatial_path=f"{base_dir}/{STAB_SPATIAL_PATH}",
+            temporal_path=f"{base_dir}/{STAB_TEMPORAL_PATH}",
+            smooth_path=f"{base_dir}/{STAB_SMOOTH_PATH}",
+            device=device,
+            mode=interpolate_mode,
+            do_linear_blend=do_linear_blend,
+        )
+
+        out_video = torch.cat(stable_list, dim=0).permute(0, 2, 3, 1)
+        out_mask = torch.cat(mask_frames_2, dim=0).permute(0, 2, 3, 1)
+        out_img2 = torch.cat(video_frames_2, dim=0)[:, [2, 1, 0], :, :].permute(
+            0, 2, 3, 1
+        )
+
+        return (
+            out_video,
+            out_img2,
+            out_mask,
+        )
+
+
+class StabStitch_Stabilize:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "video_1": ("IMAGE",),
+                # mode = 'FAST': use F.grid_sample to interpolate. It's fast, but may produce thin black boundary.
+                # mode = 'NORMAL': use our implemented interpolation function. It's a bit slower, but avoid the black boundary.
+                "interpolate_mode": (
+                    STAB_CONSTRUCT_MODES,
+                    {"default": STAB_CONSTRUCT_DEFAULT},
+                ),
+                "do_linear_blend": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+            },
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "IMAGE",
+        "IMAGE",
+    )
+    RETURN_NAMES = (
+        "combined",
+        "stab",
+        "mask",
+    )
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/StabStitch"
+
+    def todo(
+        self,
+        video_1: torch.Tensor,
+        interpolate_mode: str,
+        do_linear_blend: bool,
+    ):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        stable_list, video_frames_2, mask_frames_2 = stabstitch_single_inference(
+            video_1,
+            spatial_path=f"{base_dir}/{STAB_SPATIAL_PATH}",
+            temporal_path=f"{base_dir}/{STAB_TEMPORAL_PATH}",
+            smooth_path=f"{base_dir}/{STAB_SMOOTH_PATH}",
+            device=device,
+            mode=interpolate_mode,
+            do_linear_blend=do_linear_blend,
+        )
+
+        out_video = torch.cat(stable_list, dim=0).permute(0, 2, 3, 1)
+        out_mask = torch.cat(mask_frames_2, dim=0).permute(0, 2, 3, 1)
+        out_img2 = torch.cat(video_frames_2, dim=0)[:, [2, 1, 0], :, :].permute(
+            0, 2, 3, 1
+        )
+
+        return (
+            out_video,
+            out_img2,
+            out_mask,
+        )
+
+
+class StabStitch_Crop_Resize:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "video_1": ("IMAGE",),
+                "crop_side": (
+                    ["left", "right"],
+                    {"default": "left"},
+                ),
+                "target_h": (
+                    "INT",
+                    {"default": 360, "min": 1, "step": 1},
+                ),
+                "target_w": (
+                    "INT",
+                    {"default": 480, "min": 1, "step": 1},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("resized",)
+    FUNCTION = "todo"
+    CATEGORY = "FM_nodes/StabStitch"
+
+    def todo(
+        self,
+        video_1: torch.Tensor,
+        crop_side: str,
+        target_h: int,
+        target_w: int,
+    ):
+        resized_video = crop_or_resize(video_1, target_h, target_w, crop_side)
+
+        return (resized_video,)

workflow/stabstitch_stabilize.json

@@ -0,0 +1,458 @@
+{
+  "last_node_id": 19,
+  "last_link_id": 44,
+  "nodes": [
+    {
+      "id": 6,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1348,
+        27
+      ],
+      "size": [
+        315,
+        536.03125
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 44
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stabmask",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stabmask_00001.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 19,
+      "type": "StabStitch_Stabilize",
+      "pos": [
+        296,
+        55
+      ],
+      "size": [
+        210,
+        122
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_1",
+          "type": "IMAGE",
+          "link": 41
+        }
+      ],
+      "outputs": [
+        {
+          "name": "combined",
+          "type": "IMAGE",
+          "links": [
+            42
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "stab",
+          "type": "IMAGE",
+          "links": [
+            43
+          ],
+          "shape": 3,
+          "slot_index": 1
+        },
+        {
+          "name": "mask",
+          "type": "IMAGE",
+          "links": [
+            44
+          ],
+          "shape": 3,
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "StabStitch_Stabilize"
+      },
+      "widgets_values": [
+        "FAST",
+        true
+      ]
+    },
+    {
+      "id": 18,
+      "type": "StabStitch_Crop_Resize",
+      "pos": [
+        406,
+        229
+      ],
+      "size": [
+        210,
+        106
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_1",
+          "type": "IMAGE",
+          "link": 40
+        }
+      ],
+      "outputs": [
+        {
+          "name": "resized",
+          "type": "IMAGE",
+          "links": [
+            41
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "StabStitch_Crop_Resize"
+      },
+      "widgets_values": [
+        "left",
+        360,
+        480
+      ]
+    },
+    {
+      "id": 17,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        45,
+        8
+      ],
+      "size": [
+        235.1999969482422,
+        398.1134866403979
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            40
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadVideo"
+      },
+      "widgets_values": {
+        "video": "cat.mp4",
+        "force_rate": 0,
+        "force_size": "Disabled",
+        "custom_width": 512,
+        "custom_height": 512,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "choose video to upload": "image",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "force_rate": 0,
+            "filename": "cat.mp4",
+            "type": "input",
+            "format": "video/mp4",
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 15,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        995,
+        28
+      ],
+      "size": [
+        315,
+        536.03125
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 43
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stab",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stab_00001.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 3,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        644,
+        27
+      ],
+      "size": [
+        315,
+        536.03125
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 42
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stabstitch",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stabstitch_00007.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      40,
+      17,
+      0,
+      18,
+      0,
+      "IMAGE"
+    ],
+    [
+      41,
+      18,
+      0,
+      19,
+      0,
+      "IMAGE"
+    ],
+    [
+      42,
+      19,
+      0,
+      3,
+      0,
+      "IMAGE"
+    ],
+    [
+      43,
+      19,
+      1,
+      15,
+      0,
+      "IMAGE"
+    ],
+    [
+      44,
+      19,
+      2,
+      6,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.9090909090909091,
+      "offset": [
+        86.31161530434264,
+        196.97150602340395
+      ]
+    }
+  },
+  "version": 0.4
+}

workflow/stabstitch_stitch.json

@@ -0,0 +1,474 @@
+{
+  "last_node_id": 18,
+  "last_link_id": 49,
+  "nodes": [
+    {
+      "id": 15,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        995,
+        28
+      ],
+      "size": [
+        315,
+        473
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 48
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stab",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stab_00002.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 6,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        1348,
+        27
+      ],
+      "size": [
+        315,
+        473
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 49
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stabmask",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stabmask_00002.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 3,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        649,
+        27
+      ],
+      "size": [
+        315,
+        473
+      ],
+      "flags": {
+        "collapsed": false
+      },
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 47
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": null
+        },
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_VideoCombine"
+      },
+      "widgets_values": {
+        "frame_rate": 20,
+        "loop_count": 0,
+        "filename_prefix": "stabstitch",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 2,
+        "save_metadata": true,
+        "pingpong": false,
+        "save_output": true,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "stabstitch_00006.mp4",
+            "subfolder": "",
+            "type": "output",
+            "format": "video/h264-mp4",
+            "frame_rate": 20
+          }
+        }
+      }
+    },
+    {
+      "id": 9,
+      "type": "VHS_LoadImagesPath",
+      "pos": [
+        73,
+        218
+      ],
+      "size": [
+        226.8000030517578,
+        363.2251838887286
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            45
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null,
+          "shape": 3,
+          "slot_index": 1
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadImagesPath"
+      },
+      "widgets_values": {
+        "directory": "J:\\AI\\stabber\\video1",
+        "image_load_cap": 0,
+        "skip_first_images": 0,
+        "select_every_nth": 1,
+        "choose folder to upload": "image",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "frame_load_cap": 0,
+            "skip_first_images": 0,
+            "filename": "J:\\AI\\stabber\\video1",
+            "type": "path",
+            "format": "folder",
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 10,
+      "type": "VHS_LoadImagesPath",
+      "pos": [
+        323,
+        213
+      ],
+      "size": [
+        226.8000030517578,
+        363.2251838887286
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "type": "VHS_BatchManager",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            46
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null,
+          "shape": 3
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "links": null,
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VHS_LoadImagesPath"
+      },
+      "widgets_values": {
+        "directory": "J:\\AI\\stabber\\video2",
+        "image_load_cap": 0,
+        "skip_first_images": 0,
+        "select_every_nth": 1,
+        "choose folder to upload": "image",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "frame_load_cap": 0,
+            "skip_first_images": 0,
+            "filename": "J:\\AI\\stabber\\video2",
+            "type": "path",
+            "format": "folder",
+            "select_every_nth": 1
+          }
+        }
+      }
+    },
+    {
+      "id": 18,
+      "type": "StabStitch",
+      "pos": [
+        330,
+        37
+      ],
+      "size": [
+        215.16154450854913,
+        122
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video_1",
+          "type": "IMAGE",
+          "link": 45
+        },
+        {
+          "name": "video_2",
+          "type": "IMAGE",
+          "link": 46
+        }
+      ],
+      "outputs": [
+        {
+          "name": "combined",
+          "type": "IMAGE",
+          "links": [
+            47
+          ],
+          "shape": 3,
+          "slot_index": 0
+        },
+        {
+          "name": "stab",
+          "type": "IMAGE",
+          "links": [
+            48
+          ],
+          "shape": 3,
+          "slot_index": 1
+        },
+        {
+          "name": "mask",
+          "type": "IMAGE",
+          "links": [
+            49
+          ],
+          "shape": 3,
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "StabStitch"
+      },
+      "widgets_values": [
+        "FAST",
+        true
+      ]
+    }
+  ],
+  "links": [
+    [
+      45,
+      9,
+      0,
+      18,
+      0,
+      "IMAGE"
+    ],
+    [
+      46,
+      10,
+      0,
+      18,
+      1,
+      "IMAGE"
+    ],
+    [
+      47,
+      18,
+      0,
+      3,
+      0,
+      "IMAGE"
+    ],
+    [
+      48,
+      18,
+      1,
+      15,
+      0,
+      "IMAGE"
+    ],
+    [
+      49,
+      18,
+      2,
+      6,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        29.8384554914509,
+        163.0904108785293
+      ]
+    }
+  },
+  "version": 0.4
+}