| | 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}
|
| |
|
