|
|
from tqdm import tqdm
|
|
|
import torch
|
|
|
from torch import nn
|
|
|
|
|
|
from .optical_flow import OpticalFlow
|
|
|
from .shelf import CoTracker, CoTracker2, Tapir
|
|
|
from dot.utils.io import read_config
|
|
|
from dot.utils.torch import sample_points, sample_mask_points, get_grid
|
|
|
|
|
|
|
|
|
class PointTracker(nn.Module):
|
|
|
def __init__(self, height, width, tracker_config, tracker_path, estimator_config, estimator_path):
|
|
|
super().__init__()
|
|
|
model_args = read_config(tracker_config)
|
|
|
model_dict = {
|
|
|
"cotracker": CoTracker,
|
|
|
"cotracker2": CoTracker2,
|
|
|
"tapir": Tapir,
|
|
|
"bootstapir": Tapir
|
|
|
}
|
|
|
self.name = model_args.name
|
|
|
self.model = model_dict[model_args.name](model_args)
|
|
|
if tracker_path is not None:
|
|
|
device = next(self.model.parameters()).device
|
|
|
self.model.load_state_dict(torch.load(tracker_path, map_location=device), strict=False)
|
|
|
self.optical_flow_estimator = OpticalFlow(height, width, estimator_config, estimator_path)
|
|
|
|
|
|
def forward(self, data, mode, **kwargs):
|
|
|
if mode == "tracks_at_motion_boundaries":
|
|
|
return self.get_tracks_at_motion_boundaries(data, **kwargs)
|
|
|
elif mode == "flow_from_last_to_first_frame":
|
|
|
return self.get_flow_from_last_to_first_frame(data, **kwargs)
|
|
|
else:
|
|
|
raise ValueError(f"Unknown mode {mode}")
|
|
|
|
|
|
def get_tracks_at_motion_boundaries(self, data, num_tracks=8192, sim_tracks=2048, sample_mode="all", **kwargs):
|
|
|
video = data["video"]
|
|
|
N, S = num_tracks, sim_tracks
|
|
|
B, T, _, H, W = video.shape
|
|
|
assert N % S == 0
|
|
|
|
|
|
|
|
|
if sample_mode == "all":
|
|
|
samples_per_step = [S // T for _ in range(T)]
|
|
|
samples_per_step[0] += S - sum(samples_per_step)
|
|
|
backward_tracking = True
|
|
|
flip = False
|
|
|
elif sample_mode == "first":
|
|
|
samples_per_step = [0 for _ in range(T)]
|
|
|
samples_per_step[0] += S
|
|
|
backward_tracking = False
|
|
|
flip = False
|
|
|
elif sample_mode == "last":
|
|
|
samples_per_step = [0 for _ in range(T)]
|
|
|
samples_per_step[0] += S
|
|
|
backward_tracking = False
|
|
|
flip = True
|
|
|
else:
|
|
|
raise ValueError(f"Unknown sample mode {sample_mode}")
|
|
|
|
|
|
if flip:
|
|
|
video = video.flip(dims=[1])
|
|
|
|
|
|
|
|
|
tracks = []
|
|
|
motion_boundaries = {}
|
|
|
cache_features = True
|
|
|
for _ in tqdm(range(N // S), desc="Track batch of points", leave=False):
|
|
|
src_points = []
|
|
|
for src_step, src_samples in enumerate(samples_per_step):
|
|
|
if src_samples == 0:
|
|
|
continue
|
|
|
if not src_step in motion_boundaries:
|
|
|
tgt_step = src_step - 1 if src_step > 0 else src_step + 1
|
|
|
data = {"src_frame": video[:, src_step], "tgt_frame": video[:, tgt_step]}
|
|
|
pred = self.optical_flow_estimator(data, mode="motion_boundaries", **kwargs)
|
|
|
motion_boundaries[src_step] = pred["motion_boundaries"]
|
|
|
src_boundaries = motion_boundaries[src_step]
|
|
|
src_points.append(sample_points(src_step, src_boundaries, src_samples))
|
|
|
src_points = torch.cat(src_points, dim=1)
|
|
|
traj, vis = self.model(video, src_points, backward_tracking, cache_features)
|
|
|
tracks.append(torch.cat([traj, vis[..., None]], dim=-1))
|
|
|
cache_features = False
|
|
|
tracks = torch.cat(tracks, dim=2)
|
|
|
|
|
|
if flip:
|
|
|
tracks = tracks.flip(dims=[1])
|
|
|
|
|
|
return {"tracks": tracks}
|
|
|
|
|
|
def get_flow_from_last_to_first_frame(self, data, sim_tracks=2048, **kwargs):
|
|
|
video = data["video"]
|
|
|
video = video.flip(dims=[1])
|
|
|
src_step = 0
|
|
|
B, T, C, H, W = video.shape
|
|
|
S = sim_tracks
|
|
|
backward_tracking = False
|
|
|
cache_features = True
|
|
|
flow = get_grid(H, W, shape=[B]).cuda()
|
|
|
flow[..., 0] = flow[..., 0] * (W - 1)
|
|
|
flow[..., 1] = flow[..., 1] * (H - 1)
|
|
|
alpha = torch.zeros(B, H, W).cuda()
|
|
|
mask = torch.ones(H, W)
|
|
|
pbar = tqdm(total=H * W // S, desc="Track batch of points", leave=False)
|
|
|
while torch.any(mask):
|
|
|
points, (i, j) = sample_mask_points(src_step, mask, S)
|
|
|
idx = i * W + j
|
|
|
points = points.cuda()[None].expand(B, -1, -1)
|
|
|
|
|
|
traj, vis = self.model(video, points, backward_tracking, cache_features)
|
|
|
traj = traj[:, -1]
|
|
|
vis = vis[:, -1].float()
|
|
|
|
|
|
|
|
|
mask = mask.view(-1)
|
|
|
mask[idx] = 0
|
|
|
mask = mask.view(H, W)
|
|
|
|
|
|
|
|
|
flow = flow.view(B, -1, 2)
|
|
|
flow[:, idx] = traj - flow[:, idx]
|
|
|
flow = flow.view(B, H, W, 2)
|
|
|
|
|
|
|
|
|
alpha = alpha.view(B, -1)
|
|
|
alpha[:, idx] = vis
|
|
|
alpha = alpha.view(B, H, W)
|
|
|
|
|
|
cache_features = False
|
|
|
pbar.update(1)
|
|
|
pbar.close()
|
|
|
return {"flow": flow, "alpha": alpha}
|
|
|
|
