commit f03775f57bec5e6734b5806db61876e3cc604ac8 Author: s9roll7 Date: Wed Jun 4 19:10:21 2025 +0900 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bd13e80 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.DS_Store +*pyc +.vscode +__pycache__ +*.egg-info +*.bak +checkpoints +results +backup \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..170bc4d --- /dev/null +++ b/README.md @@ -0,0 +1,16 @@ +# Comfyui CoTracker Node + +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. + + +## Example Workflow +![example_workflow]() + +## Changelog +### 2025-6-4 +1st commit + +### Related resources +- [CoTracker](https://github.com/facebookresearch/co-tracker) +- [ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ab9b3fb --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .cotracker_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/cotracker_node.py b/cotracker_node.py new file mode 100644 index 0000000..2f95cb2 --- /dev/null +++ b/cotracker_node.py @@ -0,0 +1,435 @@ +import torch +import numpy as np +import json +import cv2 +from PIL import Image +import torchvision.transforms as transforms +import gc + +import comfy.model_management as mm + + + +class CoTrackerNode: + + def __init__(self): + self.device = mm.get_torch_device() + self.offload_device = mm.unet_offload_device() + self.model = None + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "tracking_points": ("STRING", { + "default": "", + "multiline": True, + "placeholder": "Enter x and y coordinates separated by a newline. This is optional — normally not needed, as points with large motion are selected automatically. \nExample:\n500,300\n200,250" + }), + "grid_size": ("INT", { + "default": 20, + "min": 0, + "max": 100, + "step": 1 + }), + "max_num_of_points": ("INT", { + "default": 100, + "min": 1, + "max": 10000, + "step": 1 + }), + }, + "optional": { + "tracking_mask": ("MASK", {"tooltip": "Mask for grid coordinates"}), + "confidence_threshold": ("FLOAT", { + "default": 0.90, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "min_distance": ("INT", { + "default": 30, + "min": 0, + "max": 500, + "step": 1 + }), + "force_offload": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("STRING","IMAGE") + RETURN_NAMES = ("tracking_results","image_with_results") + FUNCTION = "track_points" + CATEGORY = "tracking" + DESCRIPTION = "https://github.com/facebookresearch/co-tracker \nIf you get an OOM error, try lowering the `grid_size`." + + + def load_model(self, model_type): + try: + if self.model is None: + print(f"Loading CoTracker model: {model_type}") + self.model = torch.hub.load("facebookresearch/co-tracker", model_type).to(self.device) + self.model.to(self.device) + self.model.eval() + print("CoTracker model loaded successfully") + except Exception as e: + raise Exception(f"Failed to load CoTracker model: {str(e)}") + + def parse_tracking_points(self, tracking_points_str): + points = [] + lines = tracking_points_str.strip().split('\n') + + for line in lines: + line = line.strip() + if line and ',' in line: + try: + x, y = line.split(',') + points.append([float(x.strip()), float(y.strip())]) + except ValueError: + print(f"parse_tracking_points : Invalid point format: {line}") + continue + + return np.array(points) + + def preprocess_images(self, images): + # (B, H, W, C) -> (1, B, C, H, W) + if len(images.shape) == 4: + images = images.permute(0, 3, 1, 2) # (B, C, H, W) + images = images.unsqueeze(0) # (1, B, C, H, W) + + images = images.float() + images = images * 255 + + return images.to(self.device) + + + def prepare_query_points(self, points, video_shape): + # video_shape:(1, B, C, H, W) + + # Set points on frame 0 (specify all points on the first frame) + query_points_tensor = [] + for x, y in points: + query_points_tensor.append([0, x, y]) # frame=0, x, y + + query_points_tensor = torch.tensor(query_points_tensor, dtype=torch.float32) + + # (1, N, 3) - (batch, points, [frame, x, y]) + query_points_tensor = query_points_tensor[None].to(self.device) + + 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): + + self.load_model("cotracker3_online") + + points = self.parse_tracking_points(tracking_points) + if len(points) == 0: + print("Info : No valid points found in tracking_points") + + if tracking_mask is not None: + print(f"{tracking_mask.shape=}") + + images_np = images.cpu().numpy() + images_np = (images_np * 255).astype(np.uint8) + + video = self.preprocess_images(images) + + queries = self.prepare_query_points(points, video.shape) + + + results = [] + + if len(points) > 0: + print(f"forward - queries") + with torch.no_grad(): + self.model( + video_chunk=video, + is_first_step=True, + grid_size=0, + queries=queries, + add_support_grid=True + ) + 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, + queries=queries, + add_support_grid=True + ) # B T N 2, B T N 1 + + results, images_np = self.format_results(pred_tracks, pred_visibility, None, confidence_threshold, points, max_num_of_points, min_distance, images_np) + + print(f"{len(results)=}") + + if len(results) >= max_num_of_points: + return (results,) + + max_num_of_points -= len(results) + else: + results = [] + + 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) + print(f"{len(results2)=}") + + results = results + results2 + + + images_with_markers = torch.from_numpy(images_np) + images_with_markers = images_with_markers.float() / 255.0 + + if force_offload: + self.model.to(self.offload_device) + mm.soft_empty_cache() + gc.collect() + + 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. + + Args: + motion_sorted_indices: Indices of points sorted in descending order of motion magnitude. + tracks: Coordinate data of points across frames. + visibility: Confidence data indicating the reliability of each point.(bool) + max_points: Maximum number of points to select. + min_distance: Minimum spatial distance required between selected points. + + Returns: + selected_indices: A list of indices for the selected points. + """ + if len(motion_sorted_indices) == 0: + return [] + + selected_indices = [] + + # Compute the representative position of each point (average position over frames with high confidence) + representative_positions = {} + + for point_idx in motion_sorted_indices: + valid_frames = visibility[:, point_idx] == True + if np.any(valid_frames): + valid_positions = tracks[valid_frames, point_idx] + representative_positions[point_idx] = np.mean(valid_positions, axis=0) + else: + # Fallback: average over all frames + representative_positions[point_idx] = np.mean(tracks[:, point_idx], axis=0) + + # Select spatially dispersed points using a greedy algorithm + for candidate_idx in motion_sorted_indices: + if len(selected_indices) >= max_points: + break + + candidate_pos = representative_positions[candidate_idx] + + # Check distance to points already selected + too_close = False + for selected_idx in selected_indices: + selected_pos = representative_positions[selected_idx] + distance = np.linalg.norm(candidate_pos - selected_pos) + + if distance < min_distance: + too_close = True + break + + # Select if sufficiently far apart + if not too_close: + selected_indices.append(candidate_idx) + + return selected_indices + + + + def select_points(self, tracks, visibility, vis_threshold=0.5, max_points=9, min_distance=60): + + n_frames, n_points, _ = tracks.shape + + # 1. Confidence filtering: calculate the average confidence for each point + avg_visibility = np.mean(visibility, axis=0) + valid_points = avg_visibility >= vis_threshold + valid_indices = np.where(valid_points)[0] + + print(f"{len(valid_points)=}") + print(f"{len(valid_indices)=}") + + if len(valid_indices) == 0: + print("Warning: No points meet the confidence criteria") + return [] + + # 2. Calculate the magnitude of motion for each point (sum of movement distances across all frames) + motion_magnitudes = [] + + for point_idx in valid_indices: + total_motion = 0.0 + valid_frame_count = 0 + + for frame_idx in range(n_frames - 1): + if (visibility[frame_idx, point_idx] == True and + visibility[frame_idx + 1, point_idx] == True): + + pos1 = tracks[frame_idx, point_idx] + pos2 = tracks[frame_idx + 1, point_idx] + distance = np.linalg.norm(pos2 - pos1) + total_motion += distance + valid_frame_count += 1 + + # Normalize by the number of frames (average movement distance) + avg_motion = total_motion / max(valid_frame_count, 1) + motion_magnitudes.append(avg_motion) + + motion_magnitudes = np.array(motion_magnitudes) + + # 3. Point selection + selected_indices = [] + + if len(valid_indices) <= max_points: + selected_indices = valid_indices.tolist() + else: + # Sort points in descending order of motion magnitude + motion_sorted_indices = valid_indices[np.argsort(motion_magnitudes)[::-1]] + + high_motion_indices = self.select_diverse_points( + motion_sorted_indices, tracks, visibility, max_points=max_points-1, min_distance=min_distance + ) + selected_indices.extend(high_motion_indices) + + # Select only one point with the smallest motion (from points not yet selected) + if len(selected_indices) < max_points: + remaining_indices = [idx for idx in motion_sorted_indices if idx not in selected_indices] + if len(remaining_indices) > 0: + # Use the previous coordinates + remaining_motions = [motion_magnitudes[np.where(valid_indices == idx)[0][0]] + for idx in remaining_indices] + min_motion_idx = remaining_indices[np.argmin(remaining_motions)] + selected_indices.append(min_motion_idx) + + return selected_indices + + + def format_results(self, tracks, visibility, mask, confidence_threshold, original_points, max_points, min_distance, images_np): + # 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) + + num_frames, num_points, _ = tracks.shape + + def filter_by_mask(trs, vis, mask): + if mask is not None: + mask = mask.cpu().numpy() + if len(mask.shape) == 3 and mask.shape[0] == 1: + mask = mask[0] + + initial_coords = trs[0] # (N, 2) + + masked_indices = [] + + for n in range(initial_coords.shape[0]): + x, y = initial_coords[n] + + if (0 <= int(x) < mask.shape[1] and + 0 <= int(y) < mask.shape[0] and + mask[int(y), int(x)] > 0): + masked_indices.append(n) + + if len(masked_indices) > 0: + filtered_tracks = trs[:, masked_indices] # (T, len(masked_indices), 2) + filtered_visibility = vis[:, masked_indices] # (T, len(masked_indices)) + else: + # empty + filtered_tracks = np.empty((tracks.shape[0], 0, 2)) + filtered_visibility = np.empty((visibility.shape[0], 0)) + + return filtered_tracks, filtered_visibility + else: + return trs, vis + + + tracks, visibility = filter_by_mask(tracks, visibility, mask) + + selected_indices = self.select_points(tracks, visibility, vis_threshold=confidence_threshold, max_points=max_points, min_distance=min_distance) + + + marker_radius = 3 + marker_thickness = -1 + marker_color = (255, 0, 0) + + # Create tracking results for each point + point_results = [] + + for point_idx in selected_indices: + point_track = [] + for frame_idx in range(num_frames): + x, y = tracks[frame_idx, point_idx] + vis = visibility[frame_idx, point_idx] + + if vis == True: + point_track.append({ + "x": int(x), + "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: + 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) + + point_results += [json.dumps(point_track)] + + return point_results, images_np + +def test(): + node = CoTrackerNode() + + tracks = np.array([[(50,50),(100,50),(50,100)],[(50,50),(100,50),(50,100)],[(50,50),(100,50),(50,100)]]) + visibility = np.array([[False,True,False],[False,True,False],[True,True,False]]) + max_points = 3 + min_distance = 10 + + selected_indices = node.select_points(tracks, visibility, max_points=max_points, min_distance=min_distance) + + print(f"{selected_indices=}") + +if __name__ == '__main__': + test() + +NODE_CLASS_MAPPINGS = { + "CoTrackerNode": CoTrackerNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "CoTrackerNode": "CoTracker Point Tracking" +} +