From dd58511d4e635647e8c3a547122c0ace8343c3ce Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 10 Dec 2025 02:26:51 +0200
Subject: [PATCH] Expose some WanMove track preview drawing options

Don't need to be so huge
---
 WanMove/nodes.py      | 10 ++++++++--
 WanMove/trajectory.py | 28 ++++++++++++++++++++--------
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/WanMove/nodes.py b/WanMove/nodes.py
index b2cd35e..df5e7b0 100644
--- a/WanMove/nodes.py
+++ b/WanMove/nodes.py
@@ -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(), )
diff --git a/WanMove/trajectory.py b/WanMove/trajectory.py
index 2fd656c..e208537 100644
--- a/WanMove/trajectory.py
+++ b/WanMove/trajectory.py
@@ -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