404 lines
16 KiB
Python
404 lines
16 KiB
Python
import logging
|
|
from typing import List, Union
|
|
|
|
import torch
|
|
|
|
from ...src.coordinates import CoordinateSystem, DrawingEngine
|
|
from ...src.coordinates import coordinate_utils, drawing_utils
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RTDrawPointsNode:
|
|
"""
|
|
High-performance point drawing node optimized for real-time applications.
|
|
Draws points (circles) on images with minimal overhead.
|
|
"""
|
|
|
|
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing"
|
|
FUNCTION = "draw_points"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image_out",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"x": (
|
|
"FLOAT",
|
|
{"default": 0.5, "forceInput": True, "tooltip": "X coordinate(s) (0.0-1.0 normalized by default)"},
|
|
),
|
|
"y": (
|
|
"FLOAT",
|
|
{"default": 0.5, "forceInput": True, "tooltip": "Y coordinate(s) (0.0-1.0 normalized by default)"},
|
|
),
|
|
"radius": ("INT", {"default": 5, "min": 1, "max": 100, "tooltip": "Radius of points in pixels"}),
|
|
"color_hex": ("STRING", {"default": "#FF0000", "tooltip": "Color in #RRGGBB format"}),
|
|
},
|
|
"optional": {
|
|
"is_normalized": (
|
|
"BOOLEAN",
|
|
{"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"},
|
|
),
|
|
"batch_mapping": (
|
|
["broadcast", "one-to-one", "all-on-first"],
|
|
{"default": "broadcast", "tooltip": "How to map coords to batch images"},
|
|
),
|
|
},
|
|
}
|
|
|
|
def draw_points(
|
|
self,
|
|
image: torch.Tensor,
|
|
x: Union[float, List[float]],
|
|
y: Union[float, List[float]],
|
|
radius: int,
|
|
color_hex: str,
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
):
|
|
"""Draw points efficiently using the DrawingEngine."""
|
|
# Setup drawing engine and coordinate system
|
|
drawing_engine, space, dimensions = drawing_utils.setup_drawing_engine_and_coordinates(image, is_normalized)
|
|
|
|
return (
|
|
drawing_engine.draw_points(
|
|
image=image,
|
|
points_x=x,
|
|
points_y=y,
|
|
radius=radius,
|
|
color=color_hex,
|
|
is_normalized=is_normalized,
|
|
batch_mapping=batch_mapping,
|
|
),
|
|
)
|
|
|
|
|
|
class RTDrawLinesNode:
|
|
"""
|
|
High-performance line drawing node optimized for real-time applications.
|
|
Draws lines on images with minimal overhead.
|
|
"""
|
|
|
|
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing"
|
|
FUNCTION = "draw_lines"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image_out",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"x1": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.25,
|
|
"forceInput": True,
|
|
"tooltip": "Start X coordinate(s) (0.0-1.0 normalized by default)",
|
|
},
|
|
),
|
|
"y1": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.25,
|
|
"forceInput": True,
|
|
"tooltip": "Start Y coordinate(s) (0.0-1.0 normalized by default)",
|
|
},
|
|
),
|
|
"x2": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.75,
|
|
"forceInput": True,
|
|
"tooltip": "End X coordinate(s) (0.0-1.0 normalized by default)",
|
|
},
|
|
),
|
|
"y2": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.75,
|
|
"forceInput": True,
|
|
"tooltip": "End Y coordinate(s) (0.0-1.0 normalized by default)",
|
|
},
|
|
),
|
|
"thickness": ("INT", {"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness in pixels"}),
|
|
"color_hex": ("STRING", {"default": "#00FF00", "tooltip": "Color in #RRGGBB format"}),
|
|
},
|
|
"optional": {
|
|
"is_normalized": (
|
|
"BOOLEAN",
|
|
{"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"},
|
|
),
|
|
"batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}),
|
|
"draw_endpoints": ("BOOLEAN", {"default": False, "tooltip": "Draw circles at line endpoints"}),
|
|
"point_radius": (
|
|
"INT",
|
|
{"default": 3, "min": 1, "max": 20, "tooltip": "Radius of endpoint circles in pixels"},
|
|
),
|
|
"draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on line"}),
|
|
"label_prefix": ("STRING", {"default": "Line", "tooltip": "Prefix for label"}),
|
|
"label_values": ("STRING", {"default": "", "tooltip": "Values for label"}),
|
|
"label_position": (["Midpoint", "Start", "End"], {"default": "Midpoint"}),
|
|
"font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}),
|
|
"batch_value_mode": (["Index-based", "First value only"], {"default": "Index-based", "tooltip": "How to apply values across batch"}),
|
|
},
|
|
}
|
|
|
|
def draw_lines(
|
|
self,
|
|
image: torch.Tensor,
|
|
x1: Union[float, List[float]],
|
|
y1: Union[float, List[float]],
|
|
x2: Union[float, List[float]],
|
|
y2: Union[float, List[float]],
|
|
thickness: int,
|
|
color_hex: str,
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
draw_endpoints: bool = False,
|
|
point_radius: int = 3,
|
|
draw_label: bool = False,
|
|
label_prefix: str = "Line",
|
|
label_values: str = "",
|
|
label_position: str = "Midpoint",
|
|
font_scale: float = 0.5,
|
|
batch_value_mode: str = "Index-based",
|
|
):
|
|
"""Draw lines efficiently using the DrawingEngine."""
|
|
# Setup drawing engine and coordinate system
|
|
drawing_engine, space, dimensions = drawing_utils.setup_drawing_engine_and_coordinates(image, is_normalized)
|
|
|
|
# Parse the values list
|
|
values_list = coordinate_utils.parse_label_values(label_values)
|
|
|
|
# Setup batch drawing
|
|
batch_size, output_batch = drawing_utils.setup_batch_drawing(image)
|
|
|
|
# Standardize inputs to lists
|
|
x1_list, y1_list, x2_list, y2_list = coordinate_utils.standardize_inputs_to_lists(x1, y1, x2, y2)
|
|
|
|
# Process each batch item
|
|
for b in range(batch_size):
|
|
# Handle different batch mapping strategies for coordinates
|
|
coord_lists = [x1_list, y1_list, x2_list, y2_list]
|
|
should_skip, batch_coords = drawing_utils.handle_batch_mapping_coordinates(b, batch_mapping, coord_lists, batch_size)
|
|
|
|
if should_skip:
|
|
continue
|
|
|
|
if batch_mapping == "one-to-one":
|
|
batch_x1, batch_y1, batch_x2, batch_y2 = batch_coords
|
|
else:
|
|
# Broadcast: all coordinates on all batches
|
|
batch_x1, batch_y1, batch_x2, batch_y2 = batch_coords
|
|
|
|
# Determine the label for this batch
|
|
current_label = coordinate_utils.determine_batch_label(values_list, b, label_prefix, batch_value_mode)
|
|
|
|
# Create a single-item tensor for this batch
|
|
single_image = output_batch[b:b+1]
|
|
|
|
# Process this batch item
|
|
result = drawing_engine.draw_lines(
|
|
image=single_image,
|
|
x1=batch_x1,
|
|
y1=batch_y1,
|
|
x2=batch_x2,
|
|
y2=batch_y2,
|
|
thickness=thickness,
|
|
color=color_hex,
|
|
is_normalized=is_normalized,
|
|
batch_mapping="broadcast", # Always broadcast within the single item
|
|
draw_endpoints=draw_endpoints,
|
|
point_radius=point_radius,
|
|
draw_label=draw_label,
|
|
label_text=current_label,
|
|
label_position=label_position,
|
|
font_scale=font_scale,
|
|
)
|
|
|
|
# Update the output batch
|
|
output_batch[b:b+1] = result
|
|
|
|
return (output_batch,)
|
|
|
|
|
|
class RTDrawPolygonNode:
|
|
"""
|
|
High-performance polygon drawing node optimized for real-time applications.
|
|
Draws polygons on images with minimal overhead.
|
|
"""
|
|
|
|
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Drawing"
|
|
FUNCTION = "draw_polygon"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image_out",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"points_x": (
|
|
"FLOAT",
|
|
{
|
|
"default": [0.25, 0.75, 0.5],
|
|
"forceInput": True,
|
|
"tooltip": "List of X coordinates for polygon vertices",
|
|
},
|
|
),
|
|
"points_y": (
|
|
"FLOAT",
|
|
{
|
|
"default": [0.25, 0.25, 0.75],
|
|
"forceInput": True,
|
|
"tooltip": "List of Y coordinates for polygon vertices",
|
|
},
|
|
),
|
|
"color_hex": ("STRING", {"default": "#0000FF", "tooltip": "Color in #RRGGBB format"}),
|
|
"thickness": (
|
|
"INT",
|
|
{"default": 2, "min": 1, "max": 20, "tooltip": "Line thickness for polygon outline"},
|
|
),
|
|
},
|
|
"optional": {
|
|
"fill": ("BOOLEAN", {"default": False, "tooltip": "Whether to fill the polygon"}),
|
|
"is_normalized": (
|
|
"BOOLEAN",
|
|
{"default": True, "tooltip": "Whether coordinates are normalized (0-1) or pixel values"},
|
|
),
|
|
"batch_mapping": (["broadcast", "one-to-one", "all-on-first"], {"default": "broadcast"}),
|
|
"draw_vertices": ("BOOLEAN", {"default": False, "tooltip": "Draw points at polygon vertices"}),
|
|
"vertex_radius": ("INT", {"default": 3, "min": 1, "max": 20, "tooltip": "Radius of vertex points"}),
|
|
"draw_label": ("BOOLEAN", {"default": False, "tooltip": "Draw text label on polygon"}),
|
|
"label_prefix": ("STRING", {"default": "Polygon", "tooltip": "Prefix for label"}),
|
|
"label_values": ("STRING", {"default": "", "tooltip": "Values for label"}),
|
|
"font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}),
|
|
"batch_value_mode": (["Index-based", "First value only"], {"default": "Index-based", "tooltip": "How to apply values across batch"}),
|
|
},
|
|
}
|
|
|
|
def draw_polygon(
|
|
self,
|
|
image: torch.Tensor,
|
|
points_x: List[float],
|
|
points_y: List[float],
|
|
color_hex: str,
|
|
thickness: int,
|
|
fill: bool = False,
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
draw_vertices: bool = False,
|
|
vertex_radius: int = 3,
|
|
draw_label: bool = False,
|
|
label_prefix: str = "Polygon",
|
|
label_values: str = "",
|
|
font_scale: float = 0.5,
|
|
batch_value_mode: str = "Index-based",
|
|
):
|
|
"""Draw polygon efficiently using the DrawingEngine."""
|
|
# Setup drawing engine and coordinate system
|
|
drawing_engine, space, dimensions = drawing_utils.setup_drawing_engine_and_coordinates(image, is_normalized)
|
|
|
|
# Parse the values list
|
|
values_list = coordinate_utils.parse_label_values(label_values)
|
|
|
|
# Setup batch drawing
|
|
batch_size, output_batch = drawing_utils.setup_batch_drawing(image)
|
|
|
|
# Process each batch item
|
|
for b in range(batch_size):
|
|
# Skip processing based on mapping strategy
|
|
if batch_mapping == "all-on-first" and b > 0:
|
|
continue
|
|
|
|
# Determine the label for this batch
|
|
current_label = coordinate_utils.determine_batch_label(values_list, b, label_prefix, batch_value_mode)
|
|
|
|
# Create a single-item tensor for this batch
|
|
single_image = output_batch[b:b+1]
|
|
|
|
# Process this batch item
|
|
result = drawing_engine.draw_polygon(
|
|
image=single_image,
|
|
points_x=points_x,
|
|
points_y=points_y,
|
|
color=color_hex,
|
|
thickness=thickness,
|
|
fill=fill,
|
|
is_normalized=is_normalized,
|
|
batch_mapping="broadcast", # Always broadcast within the single item
|
|
draw_vertices=draw_vertices,
|
|
vertex_radius=vertex_radius,
|
|
draw_label=draw_label,
|
|
label_text=current_label,
|
|
font_scale=font_scale,
|
|
)
|
|
|
|
# Update the output batch
|
|
output_batch[b:b+1] = result
|
|
|
|
return (output_batch,)
|
|
|
|
|
|
class RTCoordinateConverterNode:
|
|
"""
|
|
Fast coordinate conversion between normalized (0-1) and pixel space.
|
|
Optimized for real-time applications with minimal overhead.
|
|
"""
|
|
|
|
CATEGORY = "Realtime Nodes/MediaPipe Vision/Common/Coordinates"
|
|
FUNCTION = "convert_coordinates"
|
|
RETURN_TYPES = ("FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("x_out", "y_out")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"x": ("FLOAT", {"default": 0.0, "forceInput": True, "tooltip": "X coordinate(s) to convert"}),
|
|
"y": ("FLOAT", {"default": 0.0, "forceInput": True, "tooltip": "Y coordinate(s) to convert"}),
|
|
"image_for_dimensions": ("IMAGE", {"tooltip": "Reference image for dimensions"}),
|
|
"mode": (["Pixel to Normalized", "Normalized to Pixel"],),
|
|
}
|
|
}
|
|
|
|
def convert_coordinates(
|
|
self, x: Union[float, List[float]], y: Union[float, List[float]], image_for_dimensions: torch.Tensor, mode: str
|
|
):
|
|
"""Convert coordinates using the unified coordinate system."""
|
|
# Setup coordinate system
|
|
space, dimensions = coordinate_utils.get_coordinate_space_and_dimensions(mode == "normalized", image_for_dimensions)
|
|
|
|
# Convert coordinates based on mode
|
|
if mode == "to_pixel":
|
|
x_out = CoordinateSystem.denormalize(x, dimensions[0], CoordinateSystem.PIXEL)
|
|
y_out = CoordinateSystem.denormalize(y, dimensions[1], CoordinateSystem.PIXEL)
|
|
elif mode == "to_normalized":
|
|
x_out = CoordinateSystem.normalize(x, dimensions[0], CoordinateSystem.PIXEL)
|
|
y_out = CoordinateSystem.normalize(y, dimensions[1], CoordinateSystem.PIXEL)
|
|
else:
|
|
# No conversion needed
|
|
x_out = x
|
|
y_out = y
|
|
|
|
return (x_out, y_out)
|
|
|
|
|
|
# Register nodes
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RTDrawPoints": RTDrawPointsNode,
|
|
"RTDrawLines": RTDrawLinesNode,
|
|
"RTDrawPolygon": RTDrawPolygonNode,
|
|
"RTCoordinateConverter": RTCoordinateConverterNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RTDrawPoints": "Draw Points (Fast)",
|
|
"RTDrawLines": "Draw Lines (Fast)",
|
|
"RTDrawPolygon": "Draw Polygon (Fast)",
|
|
"RTCoordinateConverter": "Coordinate Converter (Fast)",
|
|
}
|