import os
import sys
import cv2
import numpy as np
import torch
import torch.nn.functional as F
import einops
from PIL import Image
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
gmflow_dir = os.path.join(parent_dir, 'deps/gmflow')
sys.path.insert(0, gmflow_dir)
from GMFlow.gmflow.gmflow import GMFlow # noqa: E702 E402 F401
from GMFlow.utils.utils import InputPadder # noqa: E702 E402
def coords_grid(b, h, w, homogeneous=False, device=None):
y, x = torch.meshgrid(torch.arange(h), torch.arange(w)) # [H, W]
stacks = [x, y]
if homogeneous:
ones = torch.ones_like(x) # [H, W]
stacks.append(ones)
grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W]
grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W]
if device is not None:
grid = grid.to(device)
return grid
def bilinear_sample(img,
sample_coords,
mode='bilinear',
padding_mode='zeros',
return_mask=False):
# img: [B, C, H, W]
# sample_coords: [B, 2, H, W] in image scale
if sample_coords.size(1) != 2: # [B, H, W, 2]
sample_coords = sample_coords.permute(0, 3, 1, 2)
b, _, h, w = sample_coords.shape
# Normalize to [-1, 1]
x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
grid = torch.stack([x_grid, y_grid], dim=-1) # [B, H, W, 2]
img = F.grid_sample(img,
grid,
mode=mode,
padding_mode=padding_mode,
align_corners=True)
if return_mask:
mask = (x_grid >= -1) & (y_grid >= -1) & (x_grid <= 1) & (
y_grid <= 1) # [B, H, W]
return img, mask
return img
def flow_warp(feature,
flow,
mask=False,
mode='bilinear',
padding_mode='zeros'):
b, c, h, w = feature.size()
assert flow.size(1) == 2
grid = coords_grid(b, h, w).to(flow.device) + flow # [B, 2, H, W]
return bilinear_sample(feature,
grid,
mode=mode,
padding_mode=padding_mode,
return_mask=mask)
def forward_backward_consistency_check(fwd_flow,
bwd_flow,
alpha=0.01,
beta=0.5,
return_confidence=False):
# fwd_flow, bwd_flow: [B, 2, H, W]
# alpha and beta values are following UnFlow
# (https://arxiv.org/abs/1711.07837)
assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow,
dim=1) # [B, H, W]
warped_bwd_flow = flow_warp(bwd_flow, fwd_flow) # [B, 2, H, W]
warped_fwd_flow = flow_warp(fwd_flow, bwd_flow) # [B, 2, H, W]
diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1) # [B, H, W]
diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)
threshold = alpha * flow_mag + beta
if return_confidence:
# fwd_occ = diff_fwd
# bwd_occ = diff_bwd
fwd_occ = torch.exp(-diff_fwd)
bwd_occ = torch.exp(-diff_bwd)
# import ipdb; ipdb.set_trace()
# Image.fromarray((bwd_occ * 255)[0,:,:].cpu().numpy().clip(0, 255).astype(np.uint8)).save("mask.png")
else:
fwd_occ = (diff_fwd > threshold).float() # [B, H, W]
bwd_occ = (diff_bwd > threshold).float()
return fwd_occ, bwd_occ
@torch.no_grad()
def get_warped_and_mask(flow_model,
image1,
image2,
image3=None,
pixel_consistency=False,
return_confidence=False):
if image3 is None:
image3 = image1[None]
padder = InputPadder(image1.shape, padding_factor=16)
# image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
image1, image2 = padder.pad(image1[None], image2[None])
results_dict = flow_model(image1,
image2,
attn_splits_list=[2],
corr_radius_list=[-1],
prop_radius_list=[-1],
pred_bidir_flow=True)
flow_pr = results_dict['flow_preds'][-1] # [B, 2, H, W]
fwd_flow = padder.unpad(flow_pr[0]).unsqueeze(0) # [1, 2, H, W]
bwd_flow = padder.unpad(flow_pr[1]).unsqueeze(0) # [1, 2, H, W]
# results_dict_ = flow_model(image2,
# image1,
# attn_splits_list=[2],
# corr_radius_list=[-1],
# prop_radius_list=[-1],
# pred_bidir_flow=True)
# flow_pr = results_dict_['flow_preds'][-1] # [B, 2, H, W]
# fwd_flow_ = padder.unpad(flow_pr[0]).unsqueeze(0) # [1, 2, H, W]
# bwd_flow_ = padder.unpad(flow_pr[1]).unsqueeze(0) # [1, 2, H, W]
# fwd_occ_, bwd_occ_ = forward_backward_consistency_check(
# fwd_flow_, bwd_flow_, return_error=True) # [1, H, W] float
fwd_occ, bwd_occ = forward_backward_consistency_check(
fwd_flow, bwd_flow) # [1, H, W] float
if pixel_consistency:
warped_image1 = flow_warp(image1, padder.pad(bwd_flow)[0])
bwd_occ = torch.clamp(
padder.pad(bwd_occ)[0] +
(abs(image2 - warped_image1).mean(dim=1) > 255 * 0.25).float(), 0,
1)
warped_results = flow_warp(image3, bwd_flow)
if return_confidence:
fwd_err, bwd_err = forward_backward_consistency_check(
fwd_flow, bwd_flow, return_confidence=return_confidence) # [1, H, W] float
return warped_results, bwd_occ, bwd_flow, bwd_err
return warped_results, bwd_occ, bwd_flow
class FlowCalc():
def __init__(self, model_path='./weights/gmflow_sintel-0c07dcb3.pth'):
flow_model = GMFlow(
feature_channels=128,
num_scales=1,
upsample_factor=8,
num_head=1,
attention_type='swin',
ffn_dim_expansion=4,
num_transformer_layers=6,
).to('cuda')
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
weights = checkpoint['model'] if 'model' in checkpoint else checkpoint
flow_model.load_state_dict(weights, strict=False)
flow_model.eval()
self.model = flow_model
@torch.no_grad()
def get_flow(self, image1, image2, save_path=None):
if save_path is not None and os.path.exists(save_path):
bwd_flow = read_flow(save_path)
return bwd_flow
image1 = torch.from_numpy(image1).permute(2, 0, 1).float()
image2 = torch.from_numpy(image2).permute(2, 0, 1).float()
padder = InputPadder(image1.shape, padding_factor=8)
image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
results_dict = self.model(image1,
image2,
attn_splits_list=[2],
corr_radius_list=[-1],
prop_radius_list=[-1],
pred_bidir_flow=True)
flow_pr = results_dict['flow_preds'][-1] # [B, 2, H, W]
fwd_flow = padder.unpad(flow_pr[0]).unsqueeze(0) # [1, 2, H, W]
bwd_flow = padder.unpad(flow_pr[1]).unsqueeze(0) # [1, 2, H, W]
fwd_occ, bwd_occ = forward_backward_consistency_check(
fwd_flow, bwd_flow) # [1, H, W] float
if save_path is not None:
flow_np = bwd_flow.cpu().numpy()
np.save(save_path, flow_np)
mask_path = os.path.splitext(save_path)[0] + '.png'
bwd_occ = bwd_occ.cpu().permute(1, 2, 0).to(
torch.long).numpy() * 255
cv2.imwrite(mask_path, bwd_occ)
return bwd_flow
@torch.no_grad()
def get_mask(self, image1, image2, save_path=None):
if save_path is not None:
mask_path = os.path.splitext(save_path)[0] + '.png'
if os.path.exists(mask_path):
return read_mask(mask_path)
image1 = torch.from_numpy(image1).permute(2, 0, 1).float()
image2 = torch.from_numpy(image2).permute(2, 0, 1).float()
padder = InputPadder(image1.shape, padding_factor=8)
image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda())
results_dict = self.model(image1,
image2,
attn_splits_list=[2],
corr_radius_list=[-1],
prop_radius_list=[-1],
pred_bidir_flow=True)
flow_pr = results_dict['flow_preds'][-1] # [B, 2, H, W]
fwd_flow = padder.unpad(flow_pr[0]).unsqueeze(0) # [1, 2, H, W]
bwd_flow = padder.unpad(flow_pr[1]).unsqueeze(0) # [1, 2, H, W]
fwd_occ, bwd_occ = forward_backward_consistency_check(
fwd_flow, bwd_flow) # [1, H, W] float
if save_path is not None:
flow_np = bwd_flow.cpu().numpy()
np.save(save_path, flow_np)
mask_path = os.path.splitext(save_path)[0] + '.png'
bwd_occ = bwd_occ.cpu().permute(1, 2, 0).to(
torch.long).numpy() * 255
cv2.imwrite(mask_path, bwd_occ)
return bwd_occ
def warp(self, img, flow, mode='bilinear'):
expand = False
if len(img.shape) == 2:
expand = True
img = np.expand_dims(img, 2)
img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
dtype = img.dtype
img = img.to(torch.float)
res = flow_warp(img, flow, mode=mode)
res = res.to(dtype)
res = res[0].cpu().permute(1, 2, 0).numpy()
if expand:
res = res[:, :, 0]
return res
def read_flow(save_path):
flow_np = np.load(save_path)
bwd_flow = torch.from_numpy(flow_np)
return bwd_flow
def read_mask(save_path):
mask_path = os.path.splitext(save_path)[0] + '.png'
mask = cv2.imread(mask_path)
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
return mask
flow_calc = FlowCalc()