nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Expose some WanMove track preview drawing options

Browse Source 
Don't need to be so huge
This commit is contained in:
kijai
2025-12-10 02:26:51 +02:00
parent 4c97d27583
commit dd58511d4e
2 changed files with 28 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
WanMove/nodes.py
+8 -2
View File
@@ -17,6 +17,12 @@ class WanVideoWanDrawWanMoveTracks:
"images": ("IMAGE",),
"tracks": ("WANMOVETRACKS",),
},
"optional": {
"line_resolution": ("INT", {"default": 24, "min": 4, "max": 64, "step": 1, "tooltip": "Number of points to use for each line segment"}),
"circle_size": ("INT", {"default": 10, "min": 1, "max": 20, "step": 1, "tooltip": "Size of the circle to draw for each track point"}),
"opacity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Opacity of the circle to draw for each track point"}),
"line_width": ("INT", {"default": 14, "min": 1, "max": 50, "step": 1, "tooltip": "Width of the line to draw for each track"}),
}
}
RETURN_TYPES = ("IMAGE",)
@@ -24,14 +30,14 @@ class WanVideoWanDrawWanMoveTracks:
FUNCTION = "execute"
CATEGORY = "WanVideoWrapper"
def execute(self, images, tracks):
def execute(self, images, tracks, line_resolution=24, circle_size=10, opacity=0.5, line_width=14):
if tracks is None or "tracks" not in tracks:
log.warning("WanVideoWanDrawWanMoveTracks: No tracks provided.")
return (images.float().cpu(), )
track = tracks["tracks"].unsqueeze(0)
track_visibility = tracks["track_visibility"].unsqueeze(0)
images_in = images * 255.0
track_video = draw_tracks_on_video(images_in, track, track_visibility)
track_video = draw_tracks_on_video(images_in, track, track_visibility, track_frame=line_resolution, circle_size=circle_size, opacity=opacity, line_width=line_width)
track_video = torch.stack([TF.to_tensor(frame) for frame in track_video], dim=0).movedim(1, -1)
return (track_video.float().cpu(), )
WanMove/trajectory.py
+20 -8
View File
@@ -211,7 +211,7 @@ def get_video_track_video(
# Visualize functions
# --------------------------
def draw_overall_gradient_polyline_on_image(image, line_width, points, start_color):
def draw_overall_gradient_polyline_on_image(image, line_width, points, start_color, opacity=1.0):
"""
- image (Image): target image to draw on.
- line_width (int): initial line width.
@@ -246,7 +246,7 @@ def draw_overall_gradient_polyline_on_image(image, line_width, points, start_col
current_length = accumulated_length + (i / steps) * segment_length
# Alpha from fully opaque to fully transparent
alpha = int(255 * (1 - current_length / total_length))
alpha = int(255 * (1 - current_length / total_length) * opacity)
color = (*start_color, alpha)
# Interpolated coordinates
@@ -278,7 +278,7 @@ def add_weighted(rgb, track):
return Image.fromarray(blend_img.astype(np.uint8))
def draw_tracks_on_video(video, tracks, visibility=None, track_frame=24):
def draw_tracks_on_video(video, tracks, visibility=None, track_frame=24, circle_size=12, opacity=0.5, line_width=16):
color_map = [
(102, 153, 255),
(0, 255, 255),
@@ -286,8 +286,6 @@ def draw_tracks_on_video(video, tracks, visibility=None, track_frame=24):
(255, 102, 204),
(0, 255, 0)
]
circle_size = 12
line_width = 16
video = video.byte().cpu().numpy() # (81, 480, 832, 3)
tracks = tracks[0].long().detach().cpu().numpy()
@@ -312,10 +310,24 @@ def draw_tracks_on_video(video, tracks, visibility=None, track_frame=24):
tracks_coord = tracks[max(t-track_frame, 0):t+1, n]
# Draw a circle
draw = ImageDraw.Draw(frame)
draw.ellipse((track_coord[0] - circle_size, track_coord[1] - circle_size, track_coord[0] + circle_size, track_coord[1] + circle_size), fill=color_map[n % len(color_map)])
#draw = ImageDraw.Draw(frame)
#draw.ellipse((track_coord[0] - circle_size, track_coord[1] - circle_size, track_coord[0] + circle_size, track_coord[1] + circle_size), fill=color_map[n % len(color_map)])
# Draw a circle with opacity
overlay = Image.new("RGBA", frame.size, (0, 0, 0, 0))
draw_overlay = ImageDraw.Draw(overlay)
circle_color = color_map[n % len(color_map)] + (int(255 * opacity),)
draw_overlay.ellipse(
(
track_coord[0] - circle_size,
track_coord[1] - circle_size,
track_coord[0] + circle_size,
track_coord[1] + circle_size
),
fill=circle_color
)
frame = add_weighted(frame, overlay) # <-- Blend the circle overlay first
# Draw the polyline
track_image = draw_overall_gradient_polyline_on_image(frame, line_width, tracks_coord, color_map[n % len(color_map)])
track_image = draw_overall_gradient_polyline_on_image(frame, line_width, tracks_coord, color_map[n % len(color_map)], opacity=opacity)
frame = add_weighted(frame, track_image)
# Save current frame