Optimize visualization drawing

comfy_extras/nodes_wanmove.py

@@ -127,41 +127,45 @@ def replace_feature(
 
 # Visualize functions
 
-def draw_overall_gradient_polyline_on_image(image, line_width, points, start_color, opacity=1.0):
+def _draw_gradient_polyline_on_overlay(overlay, line_width, points, start_color, opacity=1.0):
-    def get_distance(p1, p2):
+    draw = ImageDraw.Draw(overlay, 'RGBA')
-        return ((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2) ** 0.5
-
-    new_image = Image.new('RGBA', image.size)
-    draw = ImageDraw.Draw(new_image, 'RGBA')
     points = points[::-1]
 
-    total_length = sum(get_distance(points[i], points[i+1]) for i in range(len(points)-1))
+    # Compute total length
+    total_length = 0
+    segment_lengths = []
+    for i in range(len(points) - 1):
+        dx = points[i + 1][0] - points[i][0]
+        dy = points[i + 1][1] - points[i][1]
+        length = (dx * dx + dy * dy) ** 0.5
+        segment_lengths.append(length)
+        total_length += length
+
+    if total_length == 0:
+        return
+
     accumulated_length = 0
 
     # Draw the gradient polyline
-    for start_point, end_point in zip(points[:-1], points[1:]):
+    for idx, (start_point, end_point) in enumerate(zip(points[:-1], points[1:])):
-        segment_length = get_distance(start_point, end_point)
+        segment_length = segment_lengths[idx]
-        steps = int(segment_length)
+        steps = max(int(segment_length), 1)
 
         for i in range(steps):
             current_length = accumulated_length + (i / steps) * segment_length
+            ratio = current_length / total_length
 
-            # Alpha from fully opaque to fully transparent
+            alpha = int(255 * (1 - ratio) * opacity)
-            alpha = int(255 * (1 - current_length / total_length) * opacity)
             color = (*start_color, alpha)
 
             x = int(start_point[0] + (end_point[0] - start_point[0]) * i / steps)
             y = int(start_point[1] + (end_point[1] - start_point[1]) * i / steps)
 
-            # Dynamic line width, decreasing from initial width to 1
+            dynamic_line_width = max(int(line_width * (1 - ratio)), 1)
-            dynamic_line_width = int(line_width * (1 - (current_length / total_length)))
-            dynamic_line_width = max(dynamic_line_width, 1)  # minimum width is 1 to avoid 0
-
             draw.line([(x, y), (x + 1, y)], fill=color, width=dynamic_line_width)
 
         accumulated_length += segment_length
 
-    return new_image
 
 def add_weighted(rgb, track):
     rgb = np.array(rgb) # [H, W, C] "RGB"
@@ -176,37 +180,60 @@ def add_weighted(rgb, track):
 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), (255, 255, 0), (255, 102, 204), (0, 255, 0)]
 
-    video = video.byte().cpu().numpy() # (81, 480, 832, 3)
+    video = video.byte().cpu().numpy()  # (81, 480, 832, 3)
     tracks = tracks[0].long().detach().cpu().numpy()
     if visibility is not None:
         visibility = visibility[0].detach().cpu().numpy()
 
-    output_frames = []
+    num_frames, height, width = video.shape[:3]
-    for t in range(video.shape[0]):
+    num_tracks = tracks.shape[1]
-        frame = video[t]
+    alpha_opacity = int(255 * opacity)
-        frame = Image.fromarray(frame).convert("RGB")
 
-        for n in range(tracks.shape[1]):
+    output_frames = []
+    for t in range(num_frames):
+        frame_rgb = video[t].astype(np.float32)
+
+        # Create a single RGBA overlay for all tracks in this frame
+        overlay = Image.new("RGBA", (width, height), (0, 0, 0, 0))
+        draw_overlay = ImageDraw.Draw(overlay)
+
+        polyline_data = []
+
+        # Draw all circles on a single overlay
+        for n in range(num_tracks):
             if visibility is not None and visibility[t, n] == 0:
                 continue
 
-            # Track coordinate at current frame
             track_coord = tracks[t, n]
-            tracks_coord = tracks[max(t-track_frame, 0):t+1, n]
+            color = color_map[n % len(color_map)]
+            circle_color = color + (alpha_opacity,)
 
-            # Draw a circle
-            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)], opacity=opacity)
-            frame = add_weighted(frame, track_image)
 
-        output_frames.append(frame.convert("RGB"))
+            # Store polyline data for batch processing
+            tracks_coord = tracks[max(t - track_frame, 0):t + 1, n]
+            if len(tracks_coord) > 1:
+                polyline_data.append((tracks_coord, color))
+
+        # Blend circles overlay once
+        overlay_np = np.array(overlay)
+        alpha = overlay_np[:, :, 3:4] / 255.0
+        frame_rgb = overlay_np[:, :, :3] * alpha + frame_rgb * (1 - alpha)
+
+        # Draw all polylines on a single overlay
+        if polyline_data:
+            polyline_overlay = Image.new("RGBA", (width, height), (0, 0, 0, 0))
+            for tracks_coord, color in polyline_data:
+                _draw_gradient_polyline_on_overlay(polyline_overlay, line_width, tracks_coord, color, opacity)
+
+            # Blend polylines overlay once
+            polyline_np = np.array(polyline_overlay)
+            alpha = polyline_np[:, :, 3:4] / 255.0
+            frame_rgb = polyline_np[:, :, :3] * alpha + frame_rgb * (1 - alpha)
+
+        output_frames.append(Image.fromarray(frame_rgb.astype(np.uint8)))
 
     return output_frames