first commit
This commit is contained in:
s9roll7
@@ -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"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user