84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Part of the implementation is borrowed and modified from DUTCode,
|
|
# publicly available at https://github.com/Annbless/DUTCode
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from tqdm import tqdm
|
|
|
|
from ..DUT.config import cfg
|
|
from .ProjectionUtils import HomoCalc, HomoProj
|
|
|
|
|
|
def mesh_warp_frame(frame, x_motion, y_motion, cap_width, cap_height):
|
|
"""
|
|
@param frame current frame [N, 1, H, W]
|
|
@param x_motion [N, 1, G_H, G_W]
|
|
@param y_motion [N, 1, G_H, G_W]
|
|
|
|
@return mesh warping according to given motion
|
|
"""
|
|
|
|
target_device = frame.device
|
|
|
|
src_grids = torch.stack(
|
|
torch.meshgrid(
|
|
torch.arange(
|
|
0, cfg.MODEL.WIDTH, cfg.MODEL.PIXELS, device=target_device),
|
|
torch.arange(
|
|
0, cfg.MODEL.HEIGHT, cfg.MODEL.PIXELS, device=target_device)),
|
|
0).permute(0, 2, 1).unsqueeze(0).float() # 2, G_H, G_W
|
|
|
|
des_grids = src_grids + torch.cat([x_motion, y_motion], 1)
|
|
|
|
projection = []
|
|
|
|
for i in range(des_grids.shape[0]):
|
|
homo = HomoCalc(src_grids[0], des_grids[i])
|
|
|
|
origin_kp = torch.stack(
|
|
torch.meshgrid(
|
|
torch.arange(0, cfg.MODEL.WIDTH, device=target_device),
|
|
torch.arange(0, cfg.MODEL.HEIGHT, device=target_device)),
|
|
0).permute(0, 2, 1).float() # 2, H, W
|
|
|
|
projected_kp = HomoProj(
|
|
homo,
|
|
origin_kp.contiguous().view(2, -1).permute(1, 0)).permute(1, 0)
|
|
|
|
projection.append(projected_kp.contiguous().view(
|
|
*origin_kp.shape).permute(1, 2, 0)) # 2, H, W --> H, W, 2
|
|
projection = torch.stack(projection, 0)
|
|
|
|
projection[:, :, :, 0] = projection[:, :, :, 0] / cfg.MODEL.WIDTH * 2. - 1.
|
|
projection[:, :, :, 1] = projection[:, :, :, 1] / \
|
|
cfg.MODEL.HEIGHT * 2. - 1.
|
|
# warp with original images
|
|
projection = projection.permute(0, 3, 1, 2)
|
|
projection = F.interpolate(
|
|
projection, (cap_height, cap_width),
|
|
mode='bilinear',
|
|
align_corners=True)
|
|
projection = projection.permute(0, 2, 3, 1)
|
|
|
|
generated_frame = F.grid_sample(frame, projection, align_corners=True)
|
|
|
|
return generated_frame
|
|
|
|
|
|
def warpListImage(images, x_motion, y_motion, cap_width, cap_height):
|
|
"""
|
|
@param images List(image [1, 1, H, W])
|
|
@param x_motion [G_H, G_W, N]
|
|
@param y_motion [G_H, G_W, N]
|
|
"""
|
|
|
|
frames = np.concatenate(images, 0)
|
|
x_motion = np.expand_dims(np.transpose(x_motion, (2, 0, 1)), 1)
|
|
y_motion = np.expand_dims(np.transpose(y_motion, (2, 0, 1)), 1)
|
|
frames = torch.from_numpy(frames.astype(np.float32))
|
|
x_motion = torch.from_numpy(x_motion.astype(np.float32))
|
|
y_motion = torch.from_numpy(y_motion.astype(np.float32))
|
|
return mesh_warp_frame(frames, x_motion, y_motion, cap_width, cap_height)
|