340 lines
13 KiB
Python
340 lines
13 KiB
Python
import logging
|
|
from typing import List, Union
|
|
|
|
import torch
|
|
|
|
from ...src.coordinates import CoordinateSystem, DrawingEngine
|
|
|
|
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."""
|
|
# Use the coordinate system to ensure proper handling
|
|
space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL
|
|
dimensions = CoordinateSystem.get_dimensions_from_tensor(image)
|
|
|
|
drawing_engine = DrawingEngine()
|
|
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_text": ("STRING", {"default": "Line", "tooltip": "Text to display as label"}),
|
|
"label_position": (["Midpoint", "Start", "End"], {"default": "Midpoint"}),
|
|
"font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}),
|
|
},
|
|
}
|
|
|
|
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_text: str = "Line",
|
|
label_position: str = "Midpoint",
|
|
font_scale: float = 0.5,
|
|
):
|
|
"""Draw lines efficiently using the DrawingEngine."""
|
|
# Use the coordinate system to ensure proper handling
|
|
space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL
|
|
dimensions = CoordinateSystem.get_dimensions_from_tensor(image)
|
|
|
|
drawing_engine = DrawingEngine()
|
|
return (
|
|
drawing_engine.draw_lines(
|
|
image=image,
|
|
x1=x1,
|
|
y1=y1,
|
|
x2=x2,
|
|
y2=y2,
|
|
thickness=thickness,
|
|
color=color_hex,
|
|
is_normalized=is_normalized,
|
|
batch_mapping=batch_mapping,
|
|
draw_endpoints=draw_endpoints,
|
|
point_radius=point_radius,
|
|
draw_label=draw_label,
|
|
label_text=label_text,
|
|
label_position=label_position,
|
|
font_scale=font_scale,
|
|
),
|
|
)
|
|
|
|
|
|
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_text": ("STRING", {"default": "Polygon", "tooltip": "Text to display as label"}),
|
|
"font_scale": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0, "step": 0.1}),
|
|
},
|
|
}
|
|
|
|
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_text: str = "Polygon",
|
|
font_scale: float = 0.5,
|
|
):
|
|
"""Draw polygon efficiently using the DrawingEngine."""
|
|
# Use the coordinate system to ensure proper handling
|
|
space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL
|
|
dimensions = CoordinateSystem.get_dimensions_from_tensor(image)
|
|
|
|
drawing_engine = DrawingEngine()
|
|
return (
|
|
drawing_engine.draw_polygon(
|
|
image=image,
|
|
points_x=points_x,
|
|
points_y=points_y,
|
|
color=color_hex,
|
|
thickness=thickness,
|
|
fill=fill,
|
|
is_normalized=is_normalized,
|
|
batch_mapping=batch_mapping,
|
|
draw_vertices=draw_vertices,
|
|
vertex_radius=vertex_radius,
|
|
draw_label=draw_label,
|
|
label_text=label_text,
|
|
label_position="Center", # Only center supported for polygons currently
|
|
font_scale=font_scale,
|
|
),
|
|
)
|
|
|
|
|
|
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 new coordinate system."""
|
|
dimensions = CoordinateSystem.get_dimensions_from_tensor(image_for_dimensions)
|
|
|
|
if mode == "Pixel to Normalized":
|
|
x_out = CoordinateSystem.normalize(x, dimensions[0], CoordinateSystem.PIXEL)
|
|
y_out = CoordinateSystem.normalize(y, dimensions[1], CoordinateSystem.PIXEL)
|
|
else: # "Normalized to Pixel"
|
|
x_out = CoordinateSystem.denormalize(x, dimensions[0], CoordinateSystem.PIXEL)
|
|
y_out = CoordinateSystem.denormalize(y, dimensions[1], CoordinateSystem.PIXEL)
|
|
|
|
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)",
|
|
}
|