diff --git a/README.md b/README.md index 6e22a7e..d42970a 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,9 @@ Amplifies coordinate movement with directional control for X/Y axes, preserving - GridPointGeneratorNode Generates a grid of coordinate points. +### 2025-6-8 +Added the enable_backward option. This is an experimental feature intended for tracking objects that don't appear in the first frame. +Fixed a bug where the min_distance option was sometimes ignored. ### Related resources diff --git a/cotracker_node.py b/cotracker_node.py index 8e1d129..144b885 100644 --- a/cotracker_node.py +++ b/cotracker_node.py @@ -7,7 +7,7 @@ import torchvision.transforms as transforms import gc import comfy.model_management as mm - +from .trajectory_integration import trajectory_integration class CoTrackerNode: @@ -57,6 +57,7 @@ class CoTrackerNode: "tooltip": "Minimum distance between tracking points" }), "force_offload": ("BOOLEAN", {"default": True}), + "enable_backward": ("BOOLEAN", {"default": False}), } } @@ -121,7 +122,7 @@ class CoTrackerNode: return query_points_tensor - def track_points(self, images, tracking_points, grid_size, max_num_of_points, tracking_mask=None, confidence_threshold=0.5, min_distance=60, force_offload=True): + def track_points(self, images, tracking_points, grid_size, max_num_of_points, tracking_mask=None, confidence_threshold=0.5, min_distance=60, force_offload=True, enable_backward=False): self.load_model("cotracker3_online") @@ -147,26 +148,31 @@ class CoTrackerNode: results = [] - if len(points) > 0: - print(f"forward - queries") + def _tracking(video, grid_size, queries, add_support_grid): with torch.no_grad(): self.model( video_chunk=video, is_first_step=True, - grid_size=0, + grid_size=grid_size, queries=queries, - add_support_grid=True + add_support_grid=add_support_grid ) for ind in range(0, video.shape[1] - self.model.step, self.model.step): pred_tracks, pred_visibility = self.model( video_chunk=video[:, ind : ind + self.model.step * 2], is_first_step=False, - grid_size=0, + grid_size=grid_size, queries=queries, - add_support_grid=True + add_support_grid=add_support_grid ) # B T N 2, B T N 1 + return pred_tracks, pred_visibility + + + if len(points) > 0: + print(f"forward - queries") - results, images_np = self.format_results(pred_tracks, pred_visibility, None, confidence_threshold, points, max_num_of_points, min_distance, images_np) + pred_tracks, pred_visibility = _tracking(video, 0, queries, True) + results, images_np = self.format_results(pred_tracks, pred_visibility, None, confidence_threshold, points, max_num_of_points, 1, images_np) print(f"{len(results)=}") @@ -179,24 +185,16 @@ class CoTrackerNode: if grid_size > 0: print(f"forward - grid") - with torch.no_grad(): - self.model( - video_chunk=video, - is_first_step=True, - grid_size=grid_size, - queries=None, - add_support_grid=False - ) - for ind in range(0, video.shape[1] - self.model.step, self.model.step): - pred_tracks, pred_visibility = self.model( - video_chunk=video[:, ind : ind + self.model.step * 2], - is_first_step=False, - grid_size=grid_size, - queries=None, - add_support_grid=False - ) # B T N 2, B T N 1 - results2, images_np = self.format_results(pred_tracks, pred_visibility, tracking_mask, confidence_threshold, points, max_num_of_points, min_distance, images_np) + pred_tracks, pred_visibility = _tracking(video, grid_size, None, False) + + if enable_backward: + pred_tracks_b, pred_visibility_b = _tracking(video.flip(1), grid_size, None, False) + _,_,_,H,W = video.shape + pred_tracks, pred_visibility = trajectory_integration(pred_tracks, pred_visibility, pred_tracks_b, pred_visibility_b, (H,W) , grid_size) + + results2, images_np = self.format_results(pred_tracks, pred_visibility, tracking_mask, confidence_threshold, points, max_num_of_points, min_distance, images_np, enable_backward) + print(f"{len(results2)=}") results = results + results2 @@ -213,7 +211,6 @@ class CoTrackerNode: return (results,images_with_markers) - def select_diverse_points(self, motion_sorted_indices, tracks, visibility, max_points, min_distance): """ Selects spatially diverse points from among those with large motion. @@ -312,7 +309,8 @@ class CoTrackerNode: # 3. Point selection selected_indices = [] - if len(valid_indices) <= max_points: +# if len(valid_indices) <= max_points: + if False: selected_indices = valid_indices.tolist() else: # Sort points in descending order of motion magnitude @@ -336,11 +334,14 @@ class CoTrackerNode: return selected_indices - def format_results(self, tracks, visibility, mask, confidence_threshold, original_points, max_points, min_distance, images_np): + def format_results(self, tracks, visibility, mask, confidence_threshold, original_points, max_points, min_distance, images_np, enable_backward=False): # tracks : (B, T, N, 2) where B=batch, T=frames, N=points tracks = tracks.squeeze(0).cpu().numpy() # (T, N, 2) visibility = visibility.squeeze(0).cpu().numpy() # (T, N) + if enable_backward: + confidence_threshold = 0 + num_frames, num_points, _ = tracks.shape def filter_by_mask(trs, vis, mask): @@ -398,17 +399,25 @@ class CoTrackerNode: "y": int(y), }) else: - # Use the previous coordinates - if len(point_track) > 0: - last_point = point_track[-1].copy() - point_track.append(last_point) - x = last_point["x"] - y = last_point["y"] - else: + if enable_backward: point_track.append({ - "x": int(x), - "y": int(y), + "x": -100, + "y": -100, }) + x = -100 + y = -100 + else: + # Use the previous coordinates + if len(point_track) > 0: + last_point = point_track[-1].copy() + point_track.append(last_point) + x = last_point["x"] + y = last_point["y"] + else: + point_track.append({ + "x": int(x), + "y": int(y), + }) if frame_idx < images_np.shape[0]: cv2.circle(images_np[frame_idx], (int(x), int(y)), marker_radius, marker_color, marker_thickness) diff --git a/pyproject.toml b/pyproject.toml index 0bfe750..868a2ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "cotracker_node" description = "This is a node that outputs tracking results of a grid or specified points using CoTracker. It can be directly connected to the WanVideo ATI Tracks Node." -version = "1.0.1" +version = "1.0.2" license = {file = "LICENSE"} [project.urls] diff --git a/trajectory_integration.py b/trajectory_integration.py new file mode 100644 index 0000000..f90c410 --- /dev/null +++ b/trajectory_integration.py @@ -0,0 +1,299 @@ +import numpy as np +import cv2 +from typing import Tuple, Optional, List +import os +import torch + +def create_mask_from_tracks(forward_tracks: np.ndarray, + forward_visibility: np.ndarray, + frame_shape: Tuple[int, int], + radius: int = 10, + frame_idx: Optional[int] = None) -> np.ndarray: + + H, W = frame_shape + + if frame_idx is not None: + mask = np.zeros((H, W), dtype=np.uint8) + points = forward_tracks[frame_idx] # shape (N, 2) + visibility = forward_visibility[frame_idx] # shape (N,) + + valid_points = points[visibility > 0] + + for point in valid_points: + x, y = int(point[0]), int(point[1]) + if 0 <= x < W and 0 <= y < H: + cv2.circle(mask, (x, y), radius, 1, -1) + + return mask + else: + T = forward_tracks.shape[0] + masks = np.zeros((T, H, W), dtype=np.uint8) + + for t in range(T): + points = forward_tracks[t] # shape (N, 2) + visibility = forward_visibility[t] # shape (N,) + + valid_points = points[visibility > 0] + + for point in valid_points: + x, y = int(point[0]), int(point[1]) + if 0 <= x < W and 0 <= y < H: + cv2.circle(masks[t], (x, y), radius, 1, -1) + + return masks + +def detect_empty_regions(forward_tracks: np.ndarray, + forward_visibility: np.ndarray, + frame_shape: Tuple[int, int], + frame_idx: int, + radius: int = 10, + min_region_size: int = 100) -> List[Tuple[int, int, int, int]]: + + mask = create_mask_from_tracks(forward_tracks, forward_visibility, frame_shape, radius, frame_idx) + + empty_mask = 1 - mask + + contours, _ = cv2.findContours(empty_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + empty_regions = [] + for contour in contours: + area = cv2.contourArea(contour) + if area >= min_region_size: + x, y, w, h = cv2.boundingRect(contour) + empty_regions.append((x, y, w, h)) + + return empty_regions + +def has_data_in_region(backward_tracks: np.ndarray, + backward_visibility: np.ndarray, + spatial_region: Tuple[int, int, int, int], + frame_idx: int, + min_points: int = 1) -> bool: + x, y, w, h = spatial_region + points = backward_tracks[frame_idx] # shape (N, 2) + visibility = backward_visibility[frame_idx] # shape (N,) + + valid_points = points[visibility > 0] + + if len(valid_points) == 0: + return False + + in_region = ((valid_points[:, 0] >= x) & (valid_points[:, 0] < x + w) & + (valid_points[:, 1] >= y) & (valid_points[:, 1] < y + h)) + + return np.sum(in_region) >= min_points + +def extract_trajectory_with_indices(tracks: np.ndarray, + visibility: np.ndarray, + spatial_region: Tuple[int, int, int, int], + frame_idx: int) -> Tuple[np.ndarray, np.ndarray, List[int]]: + + x, y, w, h = spatial_region + points = tracks[frame_idx] # shape (N, 2) + vis = visibility[frame_idx] # shape (N,) + + valid_mask = vis > 0 + in_region_mask = ((points[:, 0] >= x) & (points[:, 0] < x + w) & + (points[:, 1] >= y) & (points[:, 1] < y + h)) + + selected_indices = np.where(valid_mask & in_region_mask)[0] + + if len(selected_indices) == 0: + return np.empty((tracks.shape[0], 0, 2)), np.empty((tracks.shape[0], 0)), [] + + extracted_tracks = tracks[:, selected_indices, :] + extracted_visibility = visibility[:, selected_indices] + + return extracted_tracks, extracted_visibility, selected_indices.tolist() + + +def time_reverse(tracks: np.ndarray, visibility: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + return np.flip(tracks, axis=0), np.flip(visibility, axis=0) + + + +def integrate_tracking_results(forward_tracks: np.ndarray, + forward_visibility: np.ndarray, + backward_tracks: np.ndarray, + backward_visibility: np.ndarray, + frame_shape: Tuple[int, int], + radius: int = 10, + min_region_size: int = 100, + min_points: int = 1, + output_dir: Optional[str] = None) -> Tuple[np.ndarray, np.ndarray]: + + T = forward_tracks.shape[0] + + backward_tracks, backward_visibility = time_reverse(backward_tracks, backward_visibility) + + if output_dir is not None: + import os + os.makedirs(output_dir, exist_ok=True) + print(f"Mask images will be saved to: {output_dir}") + + integrated_tracks_list = [forward_tracks] + integrated_visibility_list = [forward_visibility] + + global_used_indices = set() + + for frame_idx in range(T): + print(f"Processing frame {frame_idx}...") + + current_integrated_tracks = np.concatenate(integrated_tracks_list, axis=1) + current_integrated_visibility = np.concatenate(integrated_visibility_list, axis=1) + + if output_dir is not None: + save_mask_visualization( + current_integrated_tracks, current_integrated_visibility, + frame_shape, frame_idx, radius, output_dir, + suffix="before_integration" + ) + + empty_regions = detect_empty_regions( + current_integrated_tracks, current_integrated_visibility, + frame_shape, frame_idx, radius, min_region_size + ) + print(f" Found {len(empty_regions)} empty regions") + + frame_new_tracks = [] + frame_new_visibility = [] + + for region_idx, spatial_region in enumerate(empty_regions): + available_indices = [i for i in range(backward_tracks.shape[1]) if i not in global_used_indices] + + if len(available_indices) == 0: + print(f" Region {region_idx}: No more available backward tracks") + continue + + available_backward_tracks = backward_tracks[:, available_indices, :] + available_backward_visibility = backward_visibility[:, available_indices] + + if has_data_in_region(available_backward_tracks, available_backward_visibility, spatial_region, frame_idx, min_points): + backward_trajectory, backward_vis, local_extracted_indices = extract_trajectory_with_indices( + available_backward_tracks, available_backward_visibility, spatial_region, frame_idx + ) + + if backward_trajectory.shape[1] > 0: + global_extracted_indices = [available_indices[i] for i in local_extracted_indices] + + print(f" Region {region_idx}: Extracted {len(global_extracted_indices)} tracks: {global_extracted_indices}") + + frame_new_tracks.append(backward_trajectory) + frame_new_visibility.append(backward_vis) + + global_used_indices.update(global_extracted_indices) + + print(f" Total used indices so far: {len(global_used_indices)}") + else: + print(f" Region {region_idx}: No backward data found") + + if frame_new_tracks: + integrated_tracks_list.extend(frame_new_tracks) + integrated_visibility_list.extend(frame_new_visibility) + + + if output_dir is not None: + final_integrated_tracks = np.concatenate(integrated_tracks_list, axis=1) + final_integrated_visibility = np.concatenate(integrated_visibility_list, axis=1) + + save_mask_visualization( + final_integrated_tracks, final_integrated_visibility, + frame_shape, frame_idx, radius, output_dir, + suffix="after_integration" + ) + + integrated_tracks = np.concatenate(integrated_tracks_list, axis=1) + integrated_visibility = np.concatenate(integrated_visibility_list, axis=1) + + + if output_dir is not None: + print(f"\nCreating comparison images...") + for frame_idx in range(T): + create_mask_comparison( + forward_tracks, forward_visibility, + integrated_tracks, integrated_visibility, + frame_shape, frame_idx, radius, output_dir + ) + + print(f"\nFinal summary:") + print(f"Used backward track indices: {sorted(global_used_indices)}") + print(f"Total backward tracks used: {len(global_used_indices)}") + print(f"Original backward tracks: {backward_tracks.shape[1]}") + print(f"Final integrated tracks shape: {integrated_tracks.shape}") + + return integrated_tracks, integrated_visibility + +def save_mask_visualization(tracks: np.ndarray, + visibility: np.ndarray, + frame_shape: Tuple[int, int], + frame_idx: int, + radius: int, + output_dir: str, + suffix: str = "") -> None: + + mask = create_mask_from_tracks(tracks, visibility, frame_shape, radius, frame_idx) + + mask_vis = (mask * 255).astype(np.uint8) + + mask_colored = cv2.applyColorMap(mask_vis, cv2.COLORMAP_JET) + + mask_colored[mask == 0] = [0, 0, 0] + + if suffix: + filename = f"mask_frame_{frame_idx:04d}_{suffix}.png" + else: + filename = f"mask_frame_{frame_idx:04d}.png" + + filepath = os.path.join(output_dir, filename) + cv2.imwrite(filepath, mask_colored) + print(f" Saved mask: {filepath}") + +def create_mask_comparison(forward_tracks: np.ndarray, + forward_visibility: np.ndarray, + integrated_tracks: np.ndarray, + integrated_visibility: np.ndarray, + frame_shape: Tuple[int, int], + frame_idx: int, + radius: int, + output_dir: str) -> None: + + forward_mask = create_mask_from_tracks(forward_tracks, forward_visibility, frame_shape, radius, frame_idx) + + integrated_mask = create_mask_from_tracks(integrated_tracks, integrated_visibility, frame_shape, radius, frame_idx) + + added_mask = integrated_mask - forward_mask + + H, W = frame_shape + comparison = np.zeros((H, W, 3), dtype=np.uint8) + + comparison[forward_mask > 0] = [255, 0, 0] + + comparison[added_mask > 0] = [0, 0, 255] + + filename = f"comparison_frame_{frame_idx:04d}.png" + filepath = os.path.join(output_dir, filename) + cv2.imwrite(filepath, comparison) + print(f" Saved comparison: {filepath}") + + +def trajectory_integration(forward_tracks, forward_visibility, backward_tracks, backward_visibility, frame_shape, grid_size): + + ft = forward_tracks.squeeze(0).cpu().numpy() # (T, N, 2) + fv = forward_visibility.squeeze(0).cpu().numpy() # (T, N) + + bt = backward_tracks.squeeze(0).cpu().numpy() # (T, N, 2) + bv = backward_visibility.squeeze(0).cpu().numpy() # (T, N) + + radius = min(frame_shape) / grid_size * 1.5 + radius = max(int(round(radius)), 3) + min_region_size = radius ** 2 + + t, v = integrate_tracking_results(ft,fv,bt,bv,frame_shape, + radius=radius, min_region_size=min_region_size, min_points=1, output_dir=None) + + t = torch.from_numpy(t).float().unsqueeze(0) + v = torch.from_numpy(v).float().unsqueeze(0) + + return t,v + +