# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import warnings import numpy as np import cv2 import torch import torch.nn.functional as F from cotracker.models.core.model_utils import smart_cat, get_points_on_a_grid from cotracker.models.build_cotracker import build_cotracker def gen_gaussian_heatmap(imgSize=200): circle_img = np.zeros((imgSize, imgSize), np.float32) circle_mask = cv2.circle(circle_img, (imgSize//2, imgSize//2), imgSize//2, 1, -1) isotropicGrayscaleImage = np.zeros((imgSize, imgSize), np.float32) # Guass Map for i in range(imgSize): for j in range(imgSize): isotropicGrayscaleImage[i, j] = 1 / 2 / np.pi / (40 ** 2) * np.exp( -1 / 2 * ((i - imgSize / 2) ** 2 / (40 ** 2) + (j - imgSize / 2) ** 2 / (40 ** 2))) isotropicGrayscaleImage = isotropicGrayscaleImage * circle_mask isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)).astype(np.float32) isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)*255).astype(np.uint8) # isotropicGrayscaleImage = cv2.resize(isotropicGrayscaleImage, (40, 40)) return isotropicGrayscaleImage def draw_heatmap(img, center_coordinate, heatmap_template, side, width, height): x1 = max(center_coordinate[0] - side, 1) x2 = min(center_coordinate[0] + side, width - 1) y1 = max(center_coordinate[1] - side, 1) y2 = min(center_coordinate[1] + side, height - 1) x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2) if (x2 - x1) < 1 or (y2 - y1) < 1: print(center_coordinate, "x1, x2, y1, y2", x1, x2, y1, y2) return img need_map = cv2.resize(heatmap_template, (x2-x1, y2-y1)) img[y1:y2,x1:x2] = need_map return img def generate_gassian_heatmap(pred_tracks, pred_visibility=None, image_size=None, side=20): width, height = image_size num_frames, num_points = pred_tracks.shape[:2] point_index_list = [point_idx for point_idx in range(num_points)] heatmap_template = gen_gaussian_heatmap() image_list = [] for frame_idx in range(num_frames): img = np.zeros((height, width), np.float32) for point_idx in point_index_list: px, py = pred_tracks[frame_idx, point_idx] if px < 0 or py < 0 or px >= width or py >= height: continue if pred_visibility is not None: if (not pred_visibility[frame_idx, point_idx]): continue img = draw_heatmap(img, (px, py), heatmap_template, side, width, height) img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_GRAY2RGB) img = torch.from_numpy(img).permute(2, 0, 1).contiguous() image_list.append(img) video_gaussion_map = torch.stack(image_list, dim=0) return video_gaussion_map # TODO: need further check and investigation def sample_trajectories( pred_tracks, pred_visibility, max_points=10, motion_threshold=1, vis_threshold=5, ): # pred_tracks: (b, f, num_points, 2) # pred_visibility: (b, f, num_points) batch_size, num_frames, num_points = pred_visibility.shape # 1. Remove points with low visibility mask = pred_visibility.sum(dim=1) > vis_threshold mask = mask.unsqueeze(1).repeat(1, num_frames, 1) pred_tracks = pred_tracks[mask].view(batch_size, num_frames, -1, 2) pred_visibility = pred_visibility[mask].view(batch_size, num_frames, -1) # 2. Thresholding: remove points with too small motions # compute the motion of each point diff = pred_tracks[:, 1:] - pred_tracks[:, :-1] # (b, f-1, num_points), sqrt(x^2 + y^2) motion = torch.norm(diff, dim=-1) # (b, num_points), mean motion for each point motion = torch.mean(motion, dim=1) # apply threshold mask = motion > motion_threshold # (b, num_points) assert mask.shape[0] == 1 num_keeped = mask.sum() if num_keeped < max_points: indices = torch.argsort(motion, dim=-1, descending=True)[:, :max_points] # (bs, max_points) mask = torch.zeros_like(mask) # (bs, num_points) # set mask to 1 for the top max_points mask[0, indices] = 1 num_keeped = mask.sum() # note sometimes mask.sum() < max_points motion = motion[mask].view(batch_size, num_keeped) # keep shape mask = mask.unsqueeze(1).repeat(1, num_frames, 1) pred_tracks = pred_tracks[mask].view(batch_size, num_frames, num_keeped, 2) pred_visibility = pred_visibility[mask].view(batch_size, num_frames, num_keeped) # 3. Sampling with larger prob for large motions num_points = min(max_points, num_keeped) if num_points == 0: warnings.warn("No points left after filtering") return None, None prob = motion / motion.max() prob = prob / prob.sum() sampled_indices = torch.multinomial(prob, num_points, replacement=False) sampled_indices = sampled_indices.squeeze(0) # (num_points, ) pred_tracks_sampled = pred_tracks[:, :, sampled_indices] pred_visibility_sampled = pred_visibility[:, :, sampled_indices] return pred_tracks_sampled, pred_visibility_sampled def sample_trajectories_with_ref( pred_tracks, pred_visibility, coords0, max_points=10, motion_threshold=1, vis_threshold=5, ): batch_size, num_frames, num_points = pred_visibility.shape visibility_sum = pred_visibility.sum(dim=1) vis_mask = visibility_sum > vis_threshold # (batch_size, num_points) pred_tracks = pred_tracks * vis_mask.unsqueeze(1).unsqueeze(-1) # (batch_size, num_frames, num_points, 2) pred_visibility = pred_visibility * vis_mask.unsqueeze(1) indices = vis_mask.nonzero(as_tuple=False) # (num_visible_points, 2) if indices.size(0) == 0: warnings.warn("No points left after visibility filtering") return None, None, None batch_indices, point_indices = indices[:, 0], indices[:, 1] coords0_filtered = coords0[batch_indices, point_indices] # (num_visible_points, 2) diff = pred_tracks[:, 1:] - pred_tracks[:, :-1] # (batch_size, num_frames-1, num_points, 2) motion = torch.norm(diff, dim=-1).mean(dim=1) # (batch_size, num_points) motion_mask = motion > motion_threshold combined_mask = vis_mask & motion_mask # (batch_size, num_points) indices = combined_mask.nonzero(as_tuple=False) if indices.size(0) == 0: warnings.warn("No points left after motion filtering") return None, None, None batch_indices, point_indices = indices[:, 0], indices[:, 1] pred_tracks_filtered = pred_tracks[batch_indices, :, point_indices, :] # (num_filtered_points, num_frames, 2) pred_visibility_filtered = pred_visibility[batch_indices, :, point_indices] # (num_filtered_points, num_frames) coords0_filtered = coords0[batch_indices, point_indices, :] # (num_filtered_points, 2) motion_filtered = motion[batch_indices, point_indices] # (num_filtered_points) num_keeped = motion_filtered.size(0) num_points_sampled = min(max_points, num_keeped) if num_points_sampled == 0: warnings.warn("No points left after filtering") return None, None, None prob = motion_filtered / motion_filtered.max() prob = prob / prob.sum() sampled_indices = torch.multinomial(prob, num_points_sampled, replacement=False) pred_tracks_sampled = pred_tracks_filtered[sampled_indices] # (num_points_sampled, num_frames, 2) pred_visibility_sampled = pred_visibility_filtered[sampled_indices] # (num_points_sampled, num_frames) coords0_sampled = coords0_filtered[sampled_indices] # (num_points_sampled, 2) pred_tracks_sampled = pred_tracks_sampled.view(batch_size, num_points_sampled, num_frames, 2).transpose(1, 2) pred_visibility_sampled = pred_visibility_sampled.view(batch_size, num_points_sampled, num_frames).transpose(1, 2) coords0_sampled = coords0_sampled.view(batch_size, num_points_sampled, 2) return pred_tracks_sampled, pred_visibility_sampled, coords0_sampled class CoTrackerPredictor(torch.nn.Module): def __init__( self, checkpoint="./checkpoints/cotracker2.pth", shift_grid=False, ): super().__init__() self.support_grid_size = 6 model = build_cotracker(checkpoint) self.interp_shape = model.model_resolution self.model = model self.model.eval() self.shift_grid = shift_grid @torch.no_grad() def forward( self, video, # (B, T, 3, H, W) # input prompt types: # - None. Dense tracks are computed in this case. You can adjust *query_frame* to compute tracks starting from a specific frame. # *backward_tracking=True* will compute tracks in both directions. # - queries. Queried points of shape (B, N, 3) in format (t, x, y) for frame index and pixel coordinates. # - grid_size. Grid of N*N points from the first frame. if segm_mask is provided, then computed only for the mask. # You can adjust *query_frame* and *backward_tracking* for the regular grid in the same way as for dense tracks. queries: torch.Tensor = None, segm_mask: torch.Tensor = None, # Segmentation mask of shape (B, 1, H, W) grid_size: int = 0, grid_query_frame: int = 0, # only for dense and regular grid tracks backward_tracking: bool = False, ): if queries is None and grid_size == 0: tracks, visibilities = self._compute_dense_tracks( video, grid_query_frame=grid_query_frame, backward_tracking=backward_tracking, ) else: tracks, visibilities = self._compute_sparse_tracks( video, queries, segm_mask, grid_size, add_support_grid=(grid_size == 0 or segm_mask is not None), grid_query_frame=grid_query_frame, backward_tracking=backward_tracking, ) return tracks, visibilities def _compute_dense_tracks(self, video, grid_query_frame, grid_size=80, backward_tracking=False): *_, H, W = video.shape grid_step = W // grid_size grid_width = W // grid_step grid_height = H // grid_step tracks = visibilities = None grid_pts = torch.zeros((1, grid_width * grid_height, 3)).to(video.device) grid_pts[0, :, 0] = grid_query_frame for offset in range(grid_step * grid_step): print(f"step {offset} / {grid_step * grid_step}") ox = offset % grid_step oy = offset // grid_step grid_pts[0, :, 1] = torch.arange(grid_width).repeat(grid_height) * grid_step + ox grid_pts[0, :, 2] = ( torch.arange(grid_height).repeat_interleave(grid_width) * grid_step + oy ) tracks_step, visibilities_step = self._compute_sparse_tracks( video=video, queries=grid_pts, backward_tracking=backward_tracking, ) tracks = smart_cat(tracks, tracks_step, dim=2) visibilities = smart_cat(visibilities, visibilities_step, dim=2) return tracks, visibilities def _compute_sparse_tracks( self, video, queries, segm_mask=None, grid_size=0, add_support_grid=False, grid_query_frame=0, backward_tracking=False, ): B, T, C, H, W = video.shape video = video.reshape(B * T, C, H, W) video = F.interpolate(video, tuple(self.interp_shape), mode="bilinear", align_corners=True) video = video.reshape(B, T, 3, self.interp_shape[0], self.interp_shape[1]) if queries is not None: B, N, D = queries.shape assert D == 3 queries = queries.clone() queries[:, :, 1:] *= queries.new_tensor( [ (self.interp_shape[1] - 1) / (W - 1), (self.interp_shape[0] - 1) / (H - 1), ] ) elif grid_size > 0: grid_pts = get_points_on_a_grid(grid_size, self.interp_shape, device=video.device, shift_grid=self.shift_grid) if segm_mask is not None: segm_mask = F.interpolate(segm_mask, tuple(self.interp_shape), mode="nearest") point_mask = segm_mask[0, 0][ (grid_pts[0, :, 1]).round().long().cpu(), (grid_pts[0, :, 0]).round().long().cpu(), ].bool() grid_pts = grid_pts[:, point_mask] queries = torch.cat( [torch.ones_like(grid_pts[:, :, :1]) * grid_query_frame, grid_pts], dim=2, ).repeat(B, 1, 1) if add_support_grid: grid_pts = get_points_on_a_grid( self.support_grid_size, self.interp_shape, device=video.device, shift_grid=self.shift_grid, ) grid_pts = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2) grid_pts = grid_pts.repeat(B, 1, 1) queries = torch.cat([queries, grid_pts], dim=1) tracks, visibilities, __ = self.model.forward(video=video, queries=queries, iters=6) if backward_tracking: tracks, visibilities = self._compute_backward_tracks( video, queries, tracks, visibilities ) if add_support_grid: queries[:, -self.support_grid_size**2 :, 0] = T - 1 if add_support_grid: tracks = tracks[:, :, : -self.support_grid_size**2] visibilities = visibilities[:, :, : -self.support_grid_size**2] thr = 0.9 visibilities = visibilities > thr # correct query-point predictions # see https://github.com/facebookresearch/co-tracker/issues/28 # TODO: batchify for i in range(len(queries)): queries_t = queries[i, : tracks.size(2), 0].to(torch.int64) arange = torch.arange(0, len(queries_t)) # overwrite the predictions with the query points tracks[i, queries_t, arange] = queries[i, : tracks.size(2), 1:] # correct visibilities, the query points should be visible visibilities[i, queries_t, arange] = True tracks *= tracks.new_tensor( [(W - 1) / (self.interp_shape[1] - 1), (H - 1) / (self.interp_shape[0] - 1)] ) return tracks, visibilities def _compute_backward_tracks(self, video, queries, tracks, visibilities): inv_video = video.flip(1).clone() inv_queries = queries.clone() inv_queries[:, :, 0] = inv_video.shape[1] - inv_queries[:, :, 0] - 1 inv_tracks, inv_visibilities, __ = self.model(video=inv_video, queries=inv_queries, iters=6) inv_tracks = inv_tracks.flip(1) inv_visibilities = inv_visibilities.flip(1) arange = torch.arange(video.shape[1], device=queries.device)[None, :, None] mask = (arange < queries[:, None, :, 0]).unsqueeze(-1).repeat(1, 1, 1, 2) tracks[mask] = inv_tracks[mask] visibilities[mask[:, :, :, 0]] = inv_visibilities[mask[:, :, :, 0]] return tracks, visibilities class CoTrackerOnlinePredictor(torch.nn.Module): def __init__(self, checkpoint="./checkpoints/cotracker2.pth"): super().__init__() self.support_grid_size = 6 model = build_cotracker(checkpoint) self.interp_shape = model.model_resolution self.step = model.window_len // 2 self.model = model self.model.eval() @torch.no_grad() def forward( self, video_chunk, is_first_step: bool = False, queries: torch.Tensor = None, grid_size: int = 10, grid_query_frame: int = 0, add_support_grid=False, ): B, T, C, H, W = video_chunk.shape # Initialize online video processing and save queried points # This needs to be done before processing *each new video* if is_first_step: self.model.init_video_online_processing() if queries is not None: B, N, D = queries.shape assert D == 3 queries = queries.clone() queries[:, :, 1:] *= queries.new_tensor( [ (self.interp_shape[1] - 1) / (W - 1), (self.interp_shape[0] - 1) / (H - 1), ] ) elif grid_size > 0: grid_pts = get_points_on_a_grid( grid_size, self.interp_shape, device=video_chunk.device ) queries = torch.cat( [torch.ones_like(grid_pts[:, :, :1]) * grid_query_frame, grid_pts], dim=2, ) if add_support_grid: grid_pts = get_points_on_a_grid( self.support_grid_size, self.interp_shape, device=video_chunk.device ) grid_pts = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2) queries = torch.cat([queries, grid_pts], dim=1) self.queries = queries return (None, None) video_chunk = video_chunk.reshape(B * T, C, H, W) video_chunk = F.interpolate( video_chunk, tuple(self.interp_shape), mode="bilinear", align_corners=True ) video_chunk = video_chunk.reshape(B, T, 3, self.interp_shape[0], self.interp_shape[1]) tracks, visibilities, __ = self.model( video=video_chunk, queries=self.queries, iters=6, is_online=True, ) thr = 0.9 return ( tracks * tracks.new_tensor( [ (W - 1) / (self.interp_shape[1] - 1), (H - 1) / (self.interp_shape[0] - 1), ] ), visibilities > thr, )