nodes/s9roll7-comfyui_cotracker_node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added the enable_backward option

Browse Source 
Fixed a bug where the min_distance option was sometimes ignored.
This commit is contained in:
s9roll7
2025-06-08 18:00:03 +09:00
parent fefc826215
commit 0618ffbd4a
4 changed files with 350 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+3
View File
@@ -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
cotracker_node.py
+47 -38
View File
@@ -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)
pyproject.toml
+1 -1
View File
@@ -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]
trajectory_integration.py
+299
View File
@@ -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