diff --git a/README.md b/README.md index 577a5f4..d6ad735 100644 --- a/README.md +++ b/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} } -``` \ No newline at end of file +``` + +``` +@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} +} +``` diff --git a/__init__.py b/__init__.py index 747eacb..8a3284d 100644 --- a/__init__.py +++ b/__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: diff --git a/constants.py b/constants.py index 2ee5fb7..0a86f06 100644 --- a/constants.py +++ b/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" \ No newline at end of file diff --git a/examples/stabstitch/video1/0001.jpg b/examples/stabstitch/video1/0001.jpg new file mode 100644 index 0000000..27a5dd3 Binary files /dev/null and b/examples/stabstitch/video1/0001.jpg differ diff --git a/examples/stabstitch/video1/0002.jpg b/examples/stabstitch/video1/0002.jpg new file mode 100644 index 0000000..50dd30e Binary files /dev/null and b/examples/stabstitch/video1/0002.jpg differ diff --git a/examples/stabstitch/video1/0003.jpg b/examples/stabstitch/video1/0003.jpg new file mode 100644 index 0000000..5afc59f Binary files /dev/null and b/examples/stabstitch/video1/0003.jpg differ diff --git a/examples/stabstitch/video1/0004.jpg b/examples/stabstitch/video1/0004.jpg new file mode 100644 index 0000000..60d1416 Binary files /dev/null and b/examples/stabstitch/video1/0004.jpg differ diff --git a/examples/stabstitch/video1/0005.jpg b/examples/stabstitch/video1/0005.jpg new file mode 100644 index 0000000..187edcf Binary files /dev/null and b/examples/stabstitch/video1/0005.jpg differ diff --git a/examples/stabstitch/video1/0006.jpg b/examples/stabstitch/video1/0006.jpg new file mode 100644 index 0000000..a72b5ac Binary files /dev/null and b/examples/stabstitch/video1/0006.jpg differ diff --git a/examples/stabstitch/video1/0007.jpg b/examples/stabstitch/video1/0007.jpg new file mode 100644 index 0000000..906e064 Binary files /dev/null and b/examples/stabstitch/video1/0007.jpg differ diff --git a/examples/stabstitch/video1/0008.jpg b/examples/stabstitch/video1/0008.jpg new file mode 100644 index 0000000..b1ecbad Binary files /dev/null and b/examples/stabstitch/video1/0008.jpg differ diff --git a/examples/stabstitch/video1/0009.jpg b/examples/stabstitch/video1/0009.jpg new file mode 100644 index 0000000..57148e7 Binary files /dev/null and b/examples/stabstitch/video1/0009.jpg differ diff --git a/examples/stabstitch/video1/0010.jpg b/examples/stabstitch/video1/0010.jpg new file mode 100644 index 0000000..e196aee Binary files /dev/null and b/examples/stabstitch/video1/0010.jpg differ diff --git a/examples/stabstitch/video1/0011.jpg b/examples/stabstitch/video1/0011.jpg new file mode 100644 index 0000000..7cc3107 Binary files /dev/null and b/examples/stabstitch/video1/0011.jpg differ diff --git a/examples/stabstitch/video1/0012.jpg b/examples/stabstitch/video1/0012.jpg new file mode 100644 index 0000000..0142920 Binary files /dev/null and b/examples/stabstitch/video1/0012.jpg differ diff --git a/examples/stabstitch/video1/0013.jpg b/examples/stabstitch/video1/0013.jpg new file mode 100644 index 0000000..05cfea7 Binary files /dev/null and b/examples/stabstitch/video1/0013.jpg differ diff --git a/examples/stabstitch/video1/0014.jpg b/examples/stabstitch/video1/0014.jpg new file mode 100644 index 0000000..bbdfeb1 Binary files /dev/null and b/examples/stabstitch/video1/0014.jpg differ diff --git a/examples/stabstitch/video1/0015.jpg b/examples/stabstitch/video1/0015.jpg new file mode 100644 index 0000000..3226876 Binary files /dev/null and b/examples/stabstitch/video1/0015.jpg differ diff --git a/examples/stabstitch/video1/0016.jpg b/examples/stabstitch/video1/0016.jpg new file mode 100644 index 0000000..1a54631 Binary files /dev/null and b/examples/stabstitch/video1/0016.jpg differ diff --git a/examples/stabstitch/video1/0017.jpg b/examples/stabstitch/video1/0017.jpg new file mode 100644 index 0000000..7353ab1 Binary files /dev/null and b/examples/stabstitch/video1/0017.jpg differ diff --git a/examples/stabstitch/video1/0018.jpg b/examples/stabstitch/video1/0018.jpg new file mode 100644 index 0000000..37e10e4 Binary files /dev/null and b/examples/stabstitch/video1/0018.jpg differ diff --git a/examples/stabstitch/video1/0019.jpg b/examples/stabstitch/video1/0019.jpg new file mode 100644 index 0000000..c91af5f Binary files /dev/null and b/examples/stabstitch/video1/0019.jpg differ diff --git a/examples/stabstitch/video1/0020.jpg b/examples/stabstitch/video1/0020.jpg new file mode 100644 index 0000000..5b84d1d Binary files /dev/null and b/examples/stabstitch/video1/0020.jpg differ diff --git a/examples/stabstitch/video2/0001.jpg b/examples/stabstitch/video2/0001.jpg new file mode 100644 index 0000000..d1a4fcf Binary files /dev/null and b/examples/stabstitch/video2/0001.jpg differ diff --git a/examples/stabstitch/video2/0002.jpg b/examples/stabstitch/video2/0002.jpg new file mode 100644 index 0000000..c5fdde7 Binary files /dev/null and b/examples/stabstitch/video2/0002.jpg differ diff --git a/examples/stabstitch/video2/0003.jpg b/examples/stabstitch/video2/0003.jpg new file mode 100644 index 0000000..9e1b481 Binary files /dev/null and b/examples/stabstitch/video2/0003.jpg differ diff --git a/examples/stabstitch/video2/0004.jpg b/examples/stabstitch/video2/0004.jpg new file mode 100644 index 0000000..a993080 Binary files /dev/null and b/examples/stabstitch/video2/0004.jpg differ diff --git a/examples/stabstitch/video2/0005.jpg b/examples/stabstitch/video2/0005.jpg new file mode 100644 index 0000000..4cd1b32 Binary files /dev/null and b/examples/stabstitch/video2/0005.jpg differ diff --git a/examples/stabstitch/video2/0006.jpg b/examples/stabstitch/video2/0006.jpg new file mode 100644 index 0000000..a41ef28 Binary files /dev/null and b/examples/stabstitch/video2/0006.jpg differ diff --git a/examples/stabstitch/video2/0007.jpg b/examples/stabstitch/video2/0007.jpg new file mode 100644 index 0000000..0a7c10a Binary files /dev/null and b/examples/stabstitch/video2/0007.jpg differ diff --git a/examples/stabstitch/video2/0008.jpg b/examples/stabstitch/video2/0008.jpg new file mode 100644 index 0000000..efe7a83 Binary files /dev/null and b/examples/stabstitch/video2/0008.jpg differ diff --git a/examples/stabstitch/video2/0009.jpg b/examples/stabstitch/video2/0009.jpg new file mode 100644 index 0000000..7a979c1 Binary files /dev/null and b/examples/stabstitch/video2/0009.jpg differ diff --git a/examples/stabstitch/video2/0010.jpg b/examples/stabstitch/video2/0010.jpg new file mode 100644 index 0000000..70d2696 Binary files /dev/null and b/examples/stabstitch/video2/0010.jpg differ diff --git a/examples/stabstitch/video2/0011.jpg b/examples/stabstitch/video2/0011.jpg new file mode 100644 index 0000000..82aa21d Binary files /dev/null and b/examples/stabstitch/video2/0011.jpg differ diff --git a/examples/stabstitch/video2/0012.jpg b/examples/stabstitch/video2/0012.jpg new file mode 100644 index 0000000..69512c4 Binary files /dev/null and b/examples/stabstitch/video2/0012.jpg differ diff --git a/examples/stabstitch/video2/0013.jpg b/examples/stabstitch/video2/0013.jpg new file mode 100644 index 0000000..d67f6e5 Binary files /dev/null and b/examples/stabstitch/video2/0013.jpg differ diff --git a/examples/stabstitch/video2/0014.jpg b/examples/stabstitch/video2/0014.jpg new file mode 100644 index 0000000..fb9f985 Binary files /dev/null and b/examples/stabstitch/video2/0014.jpg differ diff --git a/examples/stabstitch/video2/0015.jpg b/examples/stabstitch/video2/0015.jpg new file mode 100644 index 0000000..8685204 Binary files /dev/null and b/examples/stabstitch/video2/0015.jpg differ diff --git a/examples/stabstitch/video2/0016.jpg b/examples/stabstitch/video2/0016.jpg new file mode 100644 index 0000000..64f29bd Binary files /dev/null and b/examples/stabstitch/video2/0016.jpg differ diff --git a/examples/stabstitch/video2/0017.jpg b/examples/stabstitch/video2/0017.jpg new file mode 100644 index 0000000..c940bd4 Binary files /dev/null and b/examples/stabstitch/video2/0017.jpg differ diff --git a/examples/stabstitch/video2/0018.jpg b/examples/stabstitch/video2/0018.jpg new file mode 100644 index 0000000..1abe614 Binary files /dev/null and b/examples/stabstitch/video2/0018.jpg differ diff --git a/examples/stabstitch/video2/0019.jpg b/examples/stabstitch/video2/0019.jpg new file mode 100644 index 0000000..0a7bdb0 Binary files /dev/null and b/examples/stabstitch/video2/0019.jpg differ diff --git a/examples/stabstitch/video2/0020.jpg b/examples/stabstitch/video2/0020.jpg new file mode 100644 index 0000000..9f54714 Binary files /dev/null and b/examples/stabstitch/video2/0020.jpg differ diff --git a/module_stabstitch/smooth_network.py b/module_stabstitch/smooth_network.py new file mode 100644 index 0000000..940109f --- /dev/null +++ b/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 diff --git a/module_stabstitch/spatial_network.py b/module_stabstitch/spatial_network.py new file mode 100644 index 0000000..63eef49 --- /dev/null +++ b/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 diff --git a/module_stabstitch/stabstitch_model.py b/module_stabstitch/stabstitch_model.py new file mode 100644 index 0000000..452edf4 --- /dev/null +++ b/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 diff --git a/module_stabstitch/temporal_network.py b/module_stabstitch/temporal_network.py new file mode 100644 index 0000000..0154f83 --- /dev/null +++ b/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 diff --git a/module_stabstitch/torch_DLT.py b/module_stabstitch/torch_DLT.py new file mode 100644 index 0000000..7cf99e2 --- /dev/null +++ b/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 diff --git a/module_stabstitch/torch_homo_transform.py b/module_stabstitch/torch_homo_transform.py new file mode 100644 index 0000000..adf5c2a --- /dev/null +++ b/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 diff --git a/module_stabstitch/torch_tps_transform.py b/module_stabstitch/torch_tps_transform.py new file mode 100644 index 0000000..e92fb5c --- /dev/null +++ b/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 diff --git a/module_stabstitch/torch_tps_transform_point.py b/module_stabstitch/torch_tps_transform_point.py new file mode 100644 index 0000000..c5c0734 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 715019a..d5f4f82 100644 --- a/pyproject.toml +++ b/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"] diff --git a/run.py b/run.py index 6173eb1..7c21dae 100644 --- a/run.py +++ b/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 diff --git a/run_stabstitch.py b/run_stabstitch.py new file mode 100644 index 0000000..ecc0a36 --- /dev/null +++ b/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,) diff --git a/workflow/stabstitch_stabilize.json b/workflow/stabstitch_stabilize.json new file mode 100644 index 0000000..9c56384 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/workflow/stabstitch_stitch.json b/workflow/stabstitch_stitch.json new file mode 100644 index 0000000..f41e1f9 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/example_realviformer.png b/workflow_examples/example_realviformer.png similarity index 100% rename from example_realviformer.png rename to workflow_examples/example_realviformer.png diff --git a/workflow_examples/example_stabstitch_stabilize.png b/workflow_examples/example_stabstitch_stabilize.png new file mode 100644 index 0000000..dad276e Binary files /dev/null and b/workflow_examples/example_stabstitch_stabilize.png differ diff --git a/workflow_examples/example_stabstitch_stitch.png b/workflow_examples/example_stabstitch_stitch.png new file mode 100644 index 0000000..378c7a5 Binary files /dev/null and b/workflow_examples/example_stabstitch_stitch.png differ diff --git a/example_vfi_mamba.png b/workflow_examples/example_vfi_mamba.png similarity index 100% rename from example_vfi_mamba.png rename to workflow_examples/example_vfi_mamba.png