227 lines
11 KiB
Python
227 lines
11 KiB
Python
import torch
|
|
import logging
|
|
from typing import Union, List
|
|
|
|
from ...src.coordinates import CoordinateSystem
|
|
from ...src.coordinates import 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)",
|
|
} |