Merge pull request #13 from ryanontheinside/feat/coordinate-mapping
Feat/coordinate mapping
This commit is contained in:
RyanOnTheInside
@@ -6,7 +6,7 @@ from .....src.mediapipe_vision.types import BLENDSHAPES_LIST # Updated type alia
|
||||
from .....src.mediapipe_vision.utils.misc_utils import get_blendshape_categories # Utility to get all category names
|
||||
|
||||
# Import base trigger/control logic if we adapt base classes later, or implement directly for now
|
||||
from ...common.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
|
||||
from ...location_landmark.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
|
||||
from .....src.mediapipe_vision.utils.delta_utils import scale_value
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -6,7 +6,7 @@ from typing import Optional, List, Any, Tuple, Dict
|
||||
|
||||
from .....src.mediapipe_vision.types import TRANSFORM_MATRIX_LIST
|
||||
from .....src.mediapipe_vision.utils.delta_utils import scale_value
|
||||
from ...common.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
|
||||
from ...location_landmark.universal_landmark_nodes import FLOAT_EQUALITY_TOLERANCE # For trigger node comparisons
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import logging
|
||||
from ...common.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
from ...location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
# Import the landmark definitions
|
||||
from .....src.mediapipe_vision.landmark_definitions import FACE_LANDMARK_TOOLTIP
|
||||
|
||||
@@ -22,6 +22,7 @@ class FaceLandmarkPositionNode(LandmarkPositionBaseNode):
|
||||
"required": {
|
||||
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}), # Use the class variable
|
||||
"landmark_index": ("INT", {"default": 0, "min": 0, "tooltip": FACE_LANDMARK_TOOLTIP}),
|
||||
"result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the face detection to use (0=first detected face)"}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ import logging
|
||||
from typing import List, Tuple, Any, Optional
|
||||
|
||||
# Import the base class
|
||||
from ...common.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
from ...location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
# Import the specific landmark type for clarity and potential future use
|
||||
from .....src.mediapipe_vision.types import HAND_LANDMARKS
|
||||
# Import the landmark definitions
|
||||
@@ -27,6 +27,7 @@ class HandLandmarkPositionNode(LandmarkPositionBaseNode):
|
||||
"required": {
|
||||
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}),
|
||||
"landmark_index": ("INT", {"default": 8, "min": 0, "max": 20, "tooltip": HAND_LANDMARK_TOOLTIP}),
|
||||
"result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the hand detection to use (0=first detected hand)"}),
|
||||
"use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
import torch
|
||||
import logging
|
||||
from typing import Union, List
|
||||
|
||||
from ....src.mediapipe_vision.location_landmark.drawing_engine 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."""
|
||||
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."""
|
||||
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."""
|
||||
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 efficiently using DrawingEngine methods."""
|
||||
if image_for_dimensions.dim() != 4:
|
||||
logger.error("Input image must be BHWC format.")
|
||||
return ([0.0] * len(x), [0.0] * len(y)) if isinstance(x, list) else (0.0, 0.0)
|
||||
|
||||
_, height, width, _ = image_for_dimensions.shape
|
||||
drawing_engine = DrawingEngine()
|
||||
|
||||
if mode == "Pixel to Normalized":
|
||||
x_out = drawing_engine.normalize_coords(x, width)
|
||||
y_out = drawing_engine.normalize_coords(y, height)
|
||||
else: # "Normalized to Pixel"
|
||||
x_out = drawing_engine.denormalize_coords(x, width)
|
||||
y_out = drawing_engine.denormalize_coords(y, height)
|
||||
|
||||
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)",
|
||||
}
|
||||
+41
-31
@@ -7,6 +7,8 @@ import cv2
|
||||
import collections
|
||||
import typing # Add typing for Dict hint
|
||||
|
||||
#TODO: consider moving base nodes to src
|
||||
|
||||
# Assuming types are correctly imported relative to this file's location
|
||||
from ....src.mediapipe_vision.types import (
|
||||
FACE_LANDMARKS, HAND_LANDMARKS, POSE_LANDMARKS,
|
||||
@@ -28,7 +30,7 @@ class LandmarkPositionBaseNode:
|
||||
Base class for extracting position lists from specific landmark types.
|
||||
Outputs values corresponding to the input batch size, using defaults for missing items.
|
||||
Includes an 'is_valid' flag indicating successful extraction for each batch item.
|
||||
Processes only the first detected object (e.g., face, hand) per batch item.
|
||||
Processes only the detected object at the specified result_index per batch item.
|
||||
RETURN_TYPES declare the semantic type (FLOAT/BOOLEAN); the actual return value may be a list.
|
||||
"""
|
||||
FUNCTION = "extract_positions"
|
||||
@@ -42,7 +44,7 @@ class LandmarkPositionBaseNode:
|
||||
|
||||
# INPUT_TYPES are defined in subclasses
|
||||
|
||||
def extract_positions(self, landmarks: list, landmark_index: int, use_world_coordinates: bool = False) -> Tuple[List[float], List[float], List[float], List[float], List[float], List[bool]]:
|
||||
def extract_positions(self, landmarks: list, landmark_index: int, result_index: int = 0, use_world_coordinates: bool = False) -> Tuple[List[float], List[float], List[float], List[float], List[float], List[bool]]:
|
||||
"""Extracts position and properties for the specified landmark across the entire batch,
|
||||
padding with defaults and providing validity flags for missing items."""
|
||||
|
||||
@@ -69,38 +71,42 @@ class LandmarkPositionBaseNode:
|
||||
|
||||
# Check if item_list is valid and contains at least one result item
|
||||
if isinstance(item_list, list) and item_list:
|
||||
result_item = item_list[0] # Process only the first detected object
|
||||
# Get the result at the specified index if it exists
|
||||
if result_index < len(item_list):
|
||||
result_item = item_list[result_index]
|
||||
|
||||
if result_item is not None:
|
||||
try:
|
||||
# --- Determine Structure and Extract Landmark List ---
|
||||
# Case 1: result_item is the list itself (FACE_LANDMARKS)
|
||||
if isinstance(result_item, list):
|
||||
target_landmarks_list = result_item
|
||||
if use_world_coordinates and self.LANDMARKS_TYPE == "FACE_LANDMARKS":
|
||||
logger.warning(f"World coordinates requested but not available for FACE_LANDMARKS. Using image coordinates.")
|
||||
if result_item is not None:
|
||||
try:
|
||||
# --- Determine Structure and Extract Landmark List ---
|
||||
# Case 1: result_item is the list itself (FACE_LANDMARKS)
|
||||
if isinstance(result_item, list):
|
||||
target_landmarks_list = result_item
|
||||
if use_world_coordinates and self.LANDMARKS_TYPE == "FACE_LANDMARKS":
|
||||
logger.warning(f"World coordinates requested but not available for FACE_LANDMARKS. Using image coordinates.")
|
||||
|
||||
# Case 2: result_item is an object with .landmarks (HAND/POSE_LANDMARKS)
|
||||
elif hasattr(result_item, 'landmarks'):
|
||||
landmarks_attr = getattr(result_item, 'landmarks', None)
|
||||
if landmarks_attr is not None:
|
||||
use_world = (use_world_coordinates and
|
||||
self.LANDMARKS_TYPE != "FACE_LANDMARKS" and
|
||||
hasattr(result_item, 'world_landmarks') and
|
||||
getattr(result_item, 'world_landmarks', None))
|
||||
if use_world:
|
||||
target_landmarks_list = result_item.world_landmarks
|
||||
else:
|
||||
target_landmarks_list = landmarks_attr
|
||||
# else: target_landmarks_list remains None
|
||||
# Case 2: result_item is an object with .landmarks (HAND/POSE_LANDMARKS)
|
||||
elif hasattr(result_item, 'landmarks'):
|
||||
landmarks_attr = getattr(result_item, 'landmarks', None)
|
||||
if landmarks_attr is not None:
|
||||
use_world = (use_world_coordinates and
|
||||
self.LANDMARKS_TYPE != "FACE_LANDMARKS" and
|
||||
hasattr(result_item, 'world_landmarks') and
|
||||
getattr(result_item, 'world_landmarks', None))
|
||||
if use_world:
|
||||
target_landmarks_list = result_item.world_landmarks
|
||||
else:
|
||||
target_landmarks_list = landmarks_attr
|
||||
# else: target_landmarks_list remains None
|
||||
|
||||
# --- Find Landmark Point by Index ---
|
||||
if isinstance(target_landmarks_list, list) and target_landmarks_list:
|
||||
landmark_point = next((lm for lm in target_landmarks_list if lm.index == landmark_index), None)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"{self.__class__.__name__}: Unexpected error processing batch item {batch_idx}: {e}", exc_info=True)
|
||||
# Keep landmark_point as None, is_valid as False
|
||||
# --- Find Landmark Point by Index ---
|
||||
if isinstance(target_landmarks_list, list) and target_landmarks_list:
|
||||
landmark_point = next((lm for lm in target_landmarks_list if lm.index == landmark_index), None)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"{self.__class__.__name__}: Unexpected error processing batch item {batch_idx}: {e}", exc_info=True)
|
||||
# Keep landmark_point as None, is_valid as False
|
||||
else:
|
||||
logger.debug(f"{self.__class__.__name__}: Result index {result_index} exceeds available detections for batch item {batch_idx} (has {len(item_list)})")
|
||||
# else: item_list was not a list or was empty, keep defaults/invalid
|
||||
|
||||
# --- Extract Values or Use Defaults, Set Validity ---
|
||||
@@ -123,6 +129,10 @@ class LandmarkPositionBaseNode:
|
||||
pres_coords.append(current_pres)
|
||||
is_valid_list.append(is_valid)
|
||||
|
||||
# Check if each list has only one element, if so return single values instead of lists
|
||||
if len(x_coords) == 1:
|
||||
return (x_coords[0], y_coords[0], z_coords[0], vis_coords[0], pres_coords[0], is_valid_list[0])
|
||||
|
||||
return (x_coords, y_coords, z_coords, vis_coords, pres_coords, is_valid_list)
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ import logging
|
||||
from typing import List, Tuple, Any, Optional
|
||||
|
||||
# Import the base class
|
||||
from ..common.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
from ..location_landmark.universal_landmark_nodes import LandmarkPositionBaseNode
|
||||
# Import the specific landmark type for clarity and potential future use
|
||||
from ....src.mediapipe_vision.types import POSE_LANDMARKS
|
||||
# Import the landmark definitions
|
||||
@@ -26,6 +26,7 @@ class PoseLandmarkPositionNode(LandmarkPositionBaseNode):
|
||||
"required": {
|
||||
"landmarks": (cls.LANDMARKS_TYPE, {"forceInput": True}),
|
||||
"landmark_index": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": POSE_LANDMARK_TOOLTIP}),
|
||||
"result_index": ("INT", {"default": 0, "min": 0, "tooltip": "Index of the pose detection to use (0=first detected pose)"}),
|
||||
"use_world_coordinates": ("BOOLEAN", {"default": False, "tooltip": "Use world coordinates if available"}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1011,3 +1011,234 @@ class FeedbackEffect:
|
||||
|
||||
return result
|
||||
|
||||
class ToString:
|
||||
"""A node that converts any input to a string."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"any": (AlwaysEqualProxy("*"),{}), # Accept any input type
|
||||
"format_string": ("STRING", {
|
||||
"default": "{}",
|
||||
"multiline": False,
|
||||
"tooltip": "Optional format string (e.g., '{:.2f}' for 2 decimal place float)"
|
||||
})
|
||||
},
|
||||
"optional": {
|
||||
"prepend": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Text to add before the converted string"
|
||||
}),
|
||||
"append": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Text to add after the converted string"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "to_string"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def to_string(self, any, format_string="{}", prepend="", append=""):
|
||||
try:
|
||||
# Try to use the format string if provided
|
||||
result = format_string.format(any)
|
||||
except (ValueError, TypeError):
|
||||
# Fall back to simple string conversion if format fails
|
||||
result = str(any)
|
||||
|
||||
# Add prepend and append text
|
||||
result = prepend + result + append
|
||||
|
||||
return (result,)
|
||||
|
||||
class RoundNode:
|
||||
"""A node that rounds a float to a specified number of decimal places."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"value": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"tooltip": "Float value to round"
|
||||
}),
|
||||
"decimal_places": ("INT", {
|
||||
"default": 2,
|
||||
"min": 0,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"tooltip": "Number of decimal places to round to"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT",)
|
||||
RETURN_NAMES = ("rounded_value",)
|
||||
FUNCTION = "round_value"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def round_value(self, value, decimal_places):
|
||||
if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)):
|
||||
rounded = [round(v, decimal_places) for v in value]
|
||||
else:
|
||||
rounded = round(value, decimal_places)
|
||||
return (rounded,)
|
||||
|
||||
class RenormalizeFloat:
|
||||
"""Renormalizes a float value from one range to another"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"value": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"forceInput": True,
|
||||
"tooltip": "Float value to renormalize"
|
||||
}),
|
||||
"input_min": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Minimum value of the input range"
|
||||
}),
|
||||
"input_max": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Maximum value of the input range"
|
||||
}),
|
||||
"output_min": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Minimum value of the output range"
|
||||
}),
|
||||
"output_max": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Maximum value of the output range"
|
||||
}),
|
||||
"clamp": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Clamp output to the output range"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT",)
|
||||
RETURN_NAMES = ("renormalized",)
|
||||
FUNCTION = "renormalize"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def renormalize(self, value, input_min, input_max, output_min, output_max, clamp):
|
||||
"""Rescales a value from an input range to an output range"""
|
||||
# Check if value is a list, tuple, numpy array, or torch tensor
|
||||
if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)):
|
||||
# Process element-wise
|
||||
result = []
|
||||
for v in value:
|
||||
result.append(self._renormalize_single(v, input_min, input_max, output_min, output_max, clamp))
|
||||
return (result,)
|
||||
else:
|
||||
# Process single value
|
||||
result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp)
|
||||
return (result,)
|
||||
|
||||
def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp):
|
||||
"""Helper method to renormalize a single value"""
|
||||
# Handle cases where input_min equals input_max
|
||||
if abs(input_max - input_min) < 1e-6:
|
||||
return output_min
|
||||
|
||||
# Calculate normalized position in input range
|
||||
normalized = (value - input_min) / (input_max - input_min)
|
||||
|
||||
# Map to output range
|
||||
result = output_min + normalized * (output_max - output_min)
|
||||
|
||||
# Apply clamping if requested
|
||||
if clamp:
|
||||
if output_min <= output_max:
|
||||
result = max(output_min, min(output_max, result))
|
||||
else:
|
||||
result = max(output_max, min(output_min, result))
|
||||
|
||||
return result
|
||||
|
||||
class RenormalizeInt:
|
||||
"""Renormalizes an integer value from one range to another"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"value": ("INT", {
|
||||
"default": 50,
|
||||
"forceInput": True,
|
||||
"tooltip": "Integer value to renormalize"
|
||||
}),
|
||||
"input_min": ("INT", {
|
||||
"default": 0,
|
||||
"tooltip": "Minimum value of the input range"
|
||||
}),
|
||||
"input_max": ("INT", {
|
||||
"default": 100,
|
||||
"tooltip": "Maximum value of the input range"
|
||||
}),
|
||||
"output_min": ("INT", {
|
||||
"default": 0,
|
||||
"tooltip": "Minimum value of the output range"
|
||||
}),
|
||||
"output_max": ("INT", {
|
||||
"default": 100,
|
||||
"tooltip": "Maximum value of the output range"
|
||||
}),
|
||||
"clamp": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Clamp output to the output range"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("renormalized",)
|
||||
FUNCTION = "renormalize"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def renormalize(self, value, input_min, input_max, output_min, output_max, clamp):
|
||||
"""Rescales an integer value from an input range to an output range"""
|
||||
# Check if value is a list, tuple, numpy array, or torch tensor
|
||||
if isinstance(value, (list, tuple, np.ndarray, torch.Tensor)):
|
||||
# Process element-wise
|
||||
result = []
|
||||
for v in value:
|
||||
result.append(self._renormalize_single(v, input_min, input_max, output_min, output_max, clamp))
|
||||
return (result,)
|
||||
else:
|
||||
# Process single value
|
||||
result = self._renormalize_single(value, input_min, input_max, output_min, output_max, clamp)
|
||||
return (result,)
|
||||
|
||||
def _renormalize_single(self, value, input_min, input_max, output_min, output_max, clamp):
|
||||
"""Helper method to renormalize a single integer value"""
|
||||
# Handle cases where input_min equals input_max
|
||||
if input_max == input_min:
|
||||
return output_min
|
||||
|
||||
# Calculate normalized position in input range
|
||||
normalized = (value - input_min) / (input_max - input_min)
|
||||
|
||||
# Map to output range and round to integer
|
||||
result = round(output_min + normalized * (output_max - output_min))
|
||||
|
||||
# Apply clamping if requested
|
||||
if clamp:
|
||||
if output_min <= output_max:
|
||||
result = max(output_min, min(output_max, result))
|
||||
else:
|
||||
result = max(output_max, min(output_min, result))
|
||||
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,425 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import Union, List, Tuple, Dict, Optional, Any
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class DrawingEngine:
|
||||
"""High-performance drawing engine for real-time applications.
|
||||
Optimized for minimal tensor-numpy conversions and efficient batch operations.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def hex_to_bgr(hex_color: str) -> Tuple[int, int, int]:
|
||||
"""Converts hex color to BGR tuple with validation."""
|
||||
hex_color = hex_color.lstrip('#')
|
||||
if len(hex_color) != 6:
|
||||
return (0, 0, 0) # Default black
|
||||
try:
|
||||
rgb = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||||
return (rgb[2], rgb[1], rgb[0]) # RGB to BGR
|
||||
except ValueError:
|
||||
return (0, 0, 0)
|
||||
|
||||
@staticmethod
|
||||
def normalize_coords(coords: Union[float, List[float]], dim_size: int) -> Union[float, List[float]]:
|
||||
"""Normalize pixel coordinates to 0.0-1.0 range."""
|
||||
if isinstance(coords, list):
|
||||
return [float(c) / dim_size for c in coords]
|
||||
return float(coords) / dim_size
|
||||
|
||||
@staticmethod
|
||||
def denormalize_coords(coords: Union[float, List[float]], dim_size: int) -> Union[float, List[float]]:
|
||||
"""Convert normalized coordinates to pixel space."""
|
||||
if isinstance(coords, list):
|
||||
return [float(max(0, min(dim_size - 1, c * dim_size))) for c in coords]
|
||||
return float(max(0, min(dim_size - 1, coords * dim_size)))
|
||||
|
||||
@staticmethod
|
||||
def prep_batch_image(image: torch.Tensor, batch_idx: int) -> Tuple[np.ndarray, torch.device, torch.dtype]:
|
||||
"""Fast preparation of single image from batch for drawing - minimal conversion."""
|
||||
if image.dim() != 4:
|
||||
raise ValueError("Input image must be BHWC format")
|
||||
|
||||
device, dtype = image.device, image.dtype
|
||||
# Get single image, convert to float32 for processing
|
||||
img_tensor = image[batch_idx].to(torch.float32)
|
||||
# Single efficient conversion to numpy
|
||||
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
|
||||
# Ensure contiguous for CV2 operations
|
||||
return np.ascontiguousarray(img_np), device, dtype
|
||||
|
||||
@staticmethod
|
||||
def tensor_from_numpy(img_np: np.ndarray, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
||||
"""Fast conversion from numpy back to tensor with proper device/dtype."""
|
||||
return torch.from_numpy(img_np.astype(np.float32) / 255.0).to(device=device, dtype=dtype)
|
||||
|
||||
def draw_points(self,
|
||||
image: torch.Tensor,
|
||||
points_x: Union[float, List[float]],
|
||||
points_y: Union[float, List[float]],
|
||||
radius: int = 5,
|
||||
color: str = "#FF0000",
|
||||
is_normalized: bool = True,
|
||||
batch_mapping: str = "broadcast") -> torch.Tensor:
|
||||
"""
|
||||
Draw points on images efficiently.
|
||||
|
||||
Args:
|
||||
image: BHWC tensor
|
||||
points_x/y: Coordinate values (normalized 0-1 or pixel)
|
||||
radius: Point radius in pixels
|
||||
color: Hex color
|
||||
is_normalized: If True, coords are 0-1 normalized
|
||||
batch_mapping: How to map points to batch ("broadcast", "one-to-one", "all-on-first")
|
||||
|
||||
Returns:
|
||||
Modified image tensor
|
||||
"""
|
||||
batch_size, height, width, _ = image.shape
|
||||
output_batch = image.clone() # Avoid modifying input
|
||||
bgr_color = self.hex_to_bgr(color)
|
||||
|
||||
# Standardize input to lists for unified processing
|
||||
x_list = points_x if isinstance(points_x, list) else [points_x]
|
||||
y_list = points_y if isinstance(points_y, list) else [points_y]
|
||||
|
||||
# Determine batch mapping strategy
|
||||
points_len = len(x_list)
|
||||
if batch_mapping == "one-to-one" and points_len != batch_size and points_len != 1:
|
||||
logger.warning(f"One-to-one mapping requested but point list length ({points_len}) != batch size ({batch_size})")
|
||||
batch_mapping = "all-on-first"
|
||||
|
||||
# Process each image in batch
|
||||
for b in range(batch_size):
|
||||
# Skip processing based on mapping strategy
|
||||
if batch_mapping == "all-on-first" and b > 0:
|
||||
continue
|
||||
|
||||
# Get points for this batch item based on mapping strategy
|
||||
if batch_mapping == "one-to-one" and points_len == batch_size:
|
||||
# One point per image
|
||||
curr_x = [x_list[b]]
|
||||
curr_y = [y_list[b]]
|
||||
elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and points_len == 1):
|
||||
# Same point(s) on all images
|
||||
curr_x = x_list
|
||||
curr_y = y_list
|
||||
else: # "all-on-first"
|
||||
# All points on first image only
|
||||
curr_x = x_list
|
||||
curr_y = y_list
|
||||
|
||||
# Skip if no points to draw
|
||||
if not curr_x:
|
||||
continue
|
||||
|
||||
# Single efficient conversion to numpy
|
||||
img_np, device, dtype = self.prep_batch_image(output_batch, b)
|
||||
|
||||
# Draw all points for this image
|
||||
for xi, yi in zip(curr_x, curr_y):
|
||||
if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)):
|
||||
continue
|
||||
|
||||
# Convert to pixel coordinates if needed
|
||||
if is_normalized:
|
||||
px = int(max(0, min(width - 1, xi * width)))
|
||||
py = int(max(0, min(height - 1, yi * height)))
|
||||
else:
|
||||
px = int(max(0, min(width - 1, xi)))
|
||||
py = int(max(0, min(height - 1, yi)))
|
||||
|
||||
# Draw circle
|
||||
cv2.circle(img_np, (px, py), radius, bgr_color, -1)
|
||||
|
||||
# Convert back to tensor and update batch
|
||||
output_batch[b] = self.tensor_from_numpy(img_np, device, dtype)
|
||||
|
||||
return output_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 = 2,
|
||||
color: str = "#00FF00",
|
||||
is_normalized: bool = True,
|
||||
batch_mapping: str = "broadcast",
|
||||
draw_endpoints: bool = False,
|
||||
point_radius: int = 3,
|
||||
draw_label: bool = False,
|
||||
label_text: str = "",
|
||||
label_position: str = "Midpoint",
|
||||
font_scale: float = 0.5) -> torch.Tensor:
|
||||
"""
|
||||
Draw lines on images efficiently.
|
||||
|
||||
Args:
|
||||
image: BHWC tensor
|
||||
x1/y1/x2/y2: Line endpoint coordinates
|
||||
thickness: Line thickness in pixels
|
||||
color: Hex color
|
||||
is_normalized: If True, coords are 0-1 normalized
|
||||
batch_mapping: How to map lines to batch
|
||||
draw_endpoints: Whether to draw circles at endpoints
|
||||
Other params: Styling options
|
||||
|
||||
Returns:
|
||||
Modified image tensor
|
||||
"""
|
||||
batch_size, height, width, _ = image.shape
|
||||
output_batch = image.clone()
|
||||
bgr_color = self.hex_to_bgr(color)
|
||||
|
||||
# Standardize input to lists
|
||||
x1_list = x1 if isinstance(x1, list) else [x1]
|
||||
y1_list = y1 if isinstance(y1, list) else [y1]
|
||||
x2_list = x2 if isinstance(x2, list) else [x2]
|
||||
y2_list = y2 if isinstance(y2, list) else [y2]
|
||||
|
||||
# Validate input lengths
|
||||
line_len = len(x1_list)
|
||||
if not all(len(coord) == line_len for coord in [y1_list, x2_list, y2_list]):
|
||||
raise ValueError("All coordinate lists must have the same length")
|
||||
|
||||
# Determine batch mapping
|
||||
if batch_mapping == "one-to-one" and line_len != batch_size and line_len != 1:
|
||||
logger.warning(f"One-to-one mapping requested but line list length ({line_len}) != batch size ({batch_size})")
|
||||
batch_mapping = "all-on-first"
|
||||
|
||||
# Process each image in batch
|
||||
for b in range(batch_size):
|
||||
# Skip processing based on mapping strategy
|
||||
if batch_mapping == "all-on-first" and b > 0:
|
||||
continue
|
||||
|
||||
# Get lines for this batch item
|
||||
if batch_mapping == "one-to-one" and line_len == batch_size:
|
||||
indices = [b]
|
||||
elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and line_len == 1):
|
||||
indices = range(line_len)
|
||||
else: # "all-on-first"
|
||||
indices = range(line_len)
|
||||
|
||||
# Skip if no lines to draw
|
||||
if not indices:
|
||||
continue
|
||||
|
||||
# Single efficient conversion to numpy
|
||||
img_np, device, dtype = self.prep_batch_image(output_batch, b)
|
||||
|
||||
# Draw all lines for this image
|
||||
for i in indices:
|
||||
_x1, _y1 = x1_list[i], y1_list[i]
|
||||
_x2, _y2 = x2_list[i], y2_list[i]
|
||||
|
||||
if not all(isinstance(c, (float, int)) for c in [_x1, _y1, _x2, _y2]):
|
||||
continue
|
||||
|
||||
# Convert to pixel coordinates if needed
|
||||
if is_normalized:
|
||||
px1 = int(max(0, min(width - 1, _x1 * width)))
|
||||
py1 = int(max(0, min(height - 1, _y1 * height)))
|
||||
px2 = int(max(0, min(width - 1, _x2 * width)))
|
||||
py2 = int(max(0, min(height - 1, _y2 * height)))
|
||||
else:
|
||||
px1 = int(max(0, min(width - 1, _x1)))
|
||||
py1 = int(max(0, min(height - 1, _y1)))
|
||||
px2 = int(max(0, min(width - 1, _x2)))
|
||||
py2 = int(max(0, min(height - 1, _y2)))
|
||||
|
||||
# Draw line
|
||||
cv2.line(img_np, (px1, py1), (px2, py2), bgr_color, thickness)
|
||||
|
||||
# Optional features
|
||||
if draw_endpoints:
|
||||
cv2.circle(img_np, (px1, py1), point_radius, bgr_color, -1)
|
||||
cv2.circle(img_np, (px2, py2), point_radius, bgr_color, -1)
|
||||
|
||||
if draw_label and label_text:
|
||||
# Calculate label position
|
||||
if label_position == "Midpoint":
|
||||
label_x = int((px1 + px2) / 2)
|
||||
label_y = int((py1 + py2) / 2)
|
||||
elif label_position == "Start":
|
||||
label_x, label_y = px1, py1
|
||||
else: # End or fallback
|
||||
label_x, label_y = px2, py2
|
||||
|
||||
# Add offset and draw text
|
||||
label_y += 10
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
cv2.putText(img_np, label_text, (label_x, label_y),
|
||||
font, font_scale, bgr_color, 1, cv2.LINE_AA)
|
||||
|
||||
# Convert back to tensor and update batch
|
||||
output_batch[b] = self.tensor_from_numpy(img_np, device, dtype)
|
||||
|
||||
return output_batch
|
||||
|
||||
def draw_polygon(self,
|
||||
image: torch.Tensor,
|
||||
points_x: List[float],
|
||||
points_y: List[float],
|
||||
color: str = "#0000FF",
|
||||
thickness: int = 2,
|
||||
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 = "",
|
||||
label_position: str = "Center",
|
||||
font_scale: float = 0.5) -> torch.Tensor:
|
||||
"""
|
||||
Draw polygons on images efficiently.
|
||||
|
||||
Args:
|
||||
image: BHWC tensor
|
||||
points_x/y: Lists of polygon vertex coordinates
|
||||
color: Hex color
|
||||
thickness: Line thickness for polygon outline
|
||||
fill: Whether to fill the polygon
|
||||
is_normalized: If True, coords are 0-1 normalized
|
||||
batch_mapping: How to map polygons to batch
|
||||
draw_vertices: Whether to draw points at polygon vertices
|
||||
vertex_radius: Radius of vertex points
|
||||
draw_label: Whether to draw a text label
|
||||
label_text: Text to display
|
||||
label_position: Position of the label ("Center" only for now)
|
||||
font_scale: Font scale
|
||||
|
||||
Returns:
|
||||
Modified image tensor
|
||||
"""
|
||||
batch_size, height, width, _ = image.shape
|
||||
output_batch = image.clone()
|
||||
bgr_color = self.hex_to_bgr(color)
|
||||
|
||||
# Validate input
|
||||
if not isinstance(points_x, list) or not isinstance(points_y, list):
|
||||
logger.error("Polygon requires lists of points")
|
||||
return output_batch
|
||||
|
||||
if len(points_x) != len(points_y):
|
||||
logger.error("points_x and points_y must have the same length")
|
||||
return output_batch
|
||||
|
||||
if len(points_x) < 3:
|
||||
logger.error("Polygon requires at least 3 points")
|
||||
return output_batch
|
||||
|
||||
# Check if we have a list of polygon point sets or just a single polygon
|
||||
is_multi_polygon = False
|
||||
if isinstance(points_x[0], list):
|
||||
is_multi_polygon = True
|
||||
if not all(isinstance(p, list) for p in points_x + points_y):
|
||||
logger.error("If providing multiple polygons, all entries in points_x/y must be lists")
|
||||
return output_batch
|
||||
if len(points_x) != len(points_y):
|
||||
logger.error("For multiple polygons, points_x and points_y lists must have the same length")
|
||||
return output_batch
|
||||
|
||||
poly_count = len(points_x)
|
||||
# Validate each polygon has matching x,y counts
|
||||
for i in range(poly_count):
|
||||
if len(points_x[i]) != len(points_y[i]):
|
||||
logger.error(f"Polygon {i}: points_x and points_y must have the same length")
|
||||
return output_batch
|
||||
if len(points_x[i]) < 3:
|
||||
logger.error(f"Polygon {i} requires at least 3 points")
|
||||
return output_batch
|
||||
else:
|
||||
# Single polygon case - wrap in outer list for consistent processing
|
||||
points_x = [points_x]
|
||||
points_y = [points_y]
|
||||
poly_count = 1
|
||||
|
||||
# Determine batch mapping
|
||||
if batch_mapping == "one-to-one" and poly_count != batch_size and poly_count != 1:
|
||||
logger.warning(f"One-to-one mapping requested but polygon count ({poly_count}) != batch size ({batch_size})")
|
||||
batch_mapping = "all-on-first"
|
||||
|
||||
# Process each image in batch
|
||||
for b in range(batch_size):
|
||||
# Skip processing based on mapping strategy
|
||||
if batch_mapping == "all-on-first" and b > 0:
|
||||
continue
|
||||
|
||||
# Get polygons for this batch item
|
||||
if batch_mapping == "one-to-one" and poly_count == batch_size:
|
||||
poly_indices = [b]
|
||||
elif batch_mapping == "broadcast" or (batch_mapping == "one-to-one" and poly_count == 1):
|
||||
poly_indices = range(poly_count)
|
||||
else: # "all-on-first"
|
||||
poly_indices = range(poly_count)
|
||||
|
||||
# Skip if no polygons to draw
|
||||
if not poly_indices:
|
||||
continue
|
||||
|
||||
# Single efficient conversion to numpy
|
||||
img_np, device, dtype = self.prep_batch_image(output_batch, b)
|
||||
|
||||
# Draw all polygons for this image
|
||||
for i in poly_indices:
|
||||
px = points_x[i]
|
||||
py = points_y[i]
|
||||
|
||||
# Convert to pixel coordinates if needed
|
||||
if is_normalized:
|
||||
# Vectorized conversion
|
||||
px_pixels = [int(max(0, min(width - 1, x * width))) for x in px]
|
||||
py_pixels = [int(max(0, min(height - 1, y * height))) for y in py]
|
||||
else:
|
||||
# Ensure within bounds
|
||||
px_pixels = [int(max(0, min(width - 1, x))) for x in px]
|
||||
py_pixels = [int(max(0, min(height - 1, y))) for y in py]
|
||||
|
||||
# Combine x,y into points format needed by OpenCV
|
||||
points = np.array(list(zip(px_pixels, py_pixels)), np.int32)
|
||||
points = points.reshape((-1, 1, 2))
|
||||
|
||||
# Draw filled polygon
|
||||
if fill:
|
||||
cv2.fillPoly(img_np, [points], bgr_color)
|
||||
|
||||
# Draw polygon outline
|
||||
cv2.polylines(img_np, [points], True, bgr_color, thickness)
|
||||
|
||||
# Draw vertices if requested
|
||||
if draw_vertices:
|
||||
for vx, vy in zip(px_pixels, py_pixels):
|
||||
cv2.circle(img_np, (vx, vy), vertex_radius, bgr_color, -1)
|
||||
|
||||
# Draw label if requested
|
||||
if draw_label and label_text:
|
||||
# Calculate centroid for label position
|
||||
centroid_x = sum(px_pixels) // len(px_pixels)
|
||||
centroid_y = sum(py_pixels) // len(py_pixels)
|
||||
|
||||
# Add offset and draw text with outline for better visibility
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
text_size = cv2.getTextSize(label_text, font, font_scale, 1)[0]
|
||||
|
||||
# Center text
|
||||
text_x = max(0, min(width - text_size[0], centroid_x - text_size[0] // 2))
|
||||
text_y = centroid_y
|
||||
|
||||
# Draw text outline/shadow for better visibility
|
||||
cv2.putText(img_np, label_text, (text_x, text_y),
|
||||
font, font_scale, (0, 0, 0), 2, cv2.LINE_AA) # Thicker black outline
|
||||
# Draw main text
|
||||
cv2.putText(img_np, label_text, (text_x, text_y),
|
||||
font, font_scale, bgr_color, 1, cv2.LINE_AA)
|
||||
|
||||
# Convert back to tensor and update batch
|
||||
output_batch[b] = self.tensor_from_numpy(img_np, device, dtype)
|
||||
|
||||
return output_batch
|
||||
Reference in New Issue
Block a user