nodes/ryanontheinside-ComfyUI_RealtimeNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

code reduction

Browse Source
This commit is contained in:
RyanOnTheInside
2025-05-27 17:14:39 -04:00
parent 7595a6ece7
commit 3a8a798a7f
5 changed files with 306 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
node_wrappers/coordinates/drawing_nodes.py
+45 -82
View File
@@ -4,6 +4,7 @@ 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__)
@@ -58,11 +59,9 @@ class RTDrawPointsNode:
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()
# 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,
@@ -168,63 +167,38 @@ class RTDrawLinesNode:
batch_value_mode: str = "Index-based",
):
"""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)
# 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 = []
if label_values and label_values.strip():
values_list = [v.strip() for v in label_values.split(',')]
values_list = coordinate_utils.parse_label_values(label_values)
# Get the batch size
batch_size = image.shape[0]
# Process images individually to apply different labels
output_batch = image.clone()
drawing_engine = DrawingEngine()
# Setup batch drawing
batch_size, output_batch = drawing_utils.setup_batch_drawing(image)
# Standardize inputs 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]
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):
# Create a single-item tensor for this batch
single_image = output_batch[b:b+1]
# 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":
# One-to-one: each batch item gets one coordinate
if b < len(x1_list):
batch_x1 = x1_list[b]
batch_y1 = y1_list[b]
batch_x2 = x2_list[b]
batch_y2 = y2_list[b]
else:
continue # Skip if not enough coordinates
elif batch_mapping == "all-on-first" and b > 0:
continue # Skip all but the first item
batch_x1, batch_y1, batch_x2, batch_y2 = batch_coords
else:
# Broadcast: all coordinates on all batches
batch_x1 = x1_list
batch_y1 = y1_list
batch_x2 = x2_list
batch_y2 = y2_list
batch_x1, batch_y1, batch_x2, batch_y2 = batch_coords
# Determine the label for this batch
if values_list:
if batch_value_mode == "Index-based" and b < len(values_list):
# Use batch index to select value
current_value = values_list[b]
else:
# Default to first value
current_value = values_list[0]
current_label = f"{label_prefix}: {current_value}"
else:
current_label = label_prefix
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(
@@ -325,21 +299,14 @@ class RTDrawPolygonNode:
batch_value_mode: str = "Index-based",
):
"""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)
# 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 = []
if label_values and label_values.strip():
values_list = [v.strip() for v in label_values.split(',')]
values_list = coordinate_utils.parse_label_values(label_values)
# Get the batch size
batch_size = image.shape[0]
# Process images individually to apply different labels
output_batch = image.clone()
drawing_engine = DrawingEngine()
# Setup batch drawing
batch_size, output_batch = drawing_utils.setup_batch_drawing(image)
# Process each batch item
for b in range(batch_size):
@@ -347,21 +314,12 @@ class RTDrawPolygonNode:
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]
# Determine the label for this batch
if values_list:
if batch_value_mode == "Index-based" and b < len(values_list):
# Use batch index to select value
current_value = values_list[b]
else:
# Default to first value
current_value = values_list[0]
current_label = f"{label_prefix}: {current_value}"
else:
current_label = label_prefix
# Process this batch item
result = drawing_engine.draw_polygon(
image=single_image,
@@ -371,12 +329,11 @@ class RTDrawPolygonNode:
thickness=thickness,
fill=fill,
is_normalized=is_normalized,
batch_mapping=batch_mapping if batch_mapping != "one-to-one" else "broadcast",
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,
label_position="Center", # Only center supported for polygons currently
font_scale=font_scale,
)
@@ -411,16 +368,22 @@ class RTCoordinateConverterNode:
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"
"""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)
node_wrappers/coordinates/position_tracking.py
+52 -81
View File
@@ -7,6 +7,7 @@ import numpy as np
import torch
from ...src.coordinates.coordinate_delta import BaseCoordinateDelta, MAX_POSITION_HISTORY, FLOAT_EQUALITY_TOLERANCE
from ...src.coordinates import coordinate_utils
# TODO: consider moving base nodes to src
# Import from the new consolidated utilities
@@ -74,34 +75,31 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta):
clamp = kwargs.get("clamp")
fallback_value = int(output_min) if self.IS_INT_OUTPUT else float(output_min)
if None in [x_in, y_in, z_in, delta_min, delta_max, output_min, output_max, clamp]:
logger.error(f"{self.__class__.__name__}: Missing required arguments.")
# Check required arguments
required_args = {
"x": x_in, "y": y_in, "z": z_in, "delta_min": delta_min,
"delta_max": delta_max, "output_min": output_min, "output_max": output_max, "clamp": clamp
}
if not coordinate_utils.check_required_arguments(required_args, self.__class__.__name__):
return (fallback_value,)
# --- Input Type Handling & History Cleanup ---
is_list = isinstance(x_in, list)
if is_list != isinstance(y_in, list) or is_list != isinstance(z_in, list):
raise ValueError(f"{self.__class__.__name__}: Inputs x, y, z must be all floats or all lists.")
# Validate input types
is_list = coordinate_utils.validate_coordinate_input_types(x_in, y_in, z_in, self.__class__.__name__)
if is_list:
current_batch_size = len(x_in)
# Remove histories for indices no longer present
keys_to_remove = [k for k in self._histories if k >= current_batch_size]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Clean up histories
coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__)
if current_batch_size == 0:
logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.")
# Handle empty lists
if coordinate_utils.handle_empty_coordinate_lists(current_batch_size, self.__class__.__name__):
return ([],)
# Process List Input
results = []
for i in range(current_batch_size):
xi, yi, zi = x_in[i], y_in[i], z_in[i]
if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]):
logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.")
if not coordinate_utils.validate_coordinate_types_in_list([[xi], [yi], [zi]], 0, self.__class__.__name__):
results.append(fallback_value)
continue
@@ -116,16 +114,11 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta):
return (results,)
else:
# Single Float Input: Remove all histories except for index 0
keys_to_remove = [k for k in self._histories if k != 0]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Clean up single input histories
coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__)
# Process Single Float Input
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.")
# Validate single coordinate types
if not coordinate_utils.validate_single_coordinate_types(x_in, y_in, z_in, self.__class__.__name__):
return (fallback_value,)
raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0)
@@ -187,34 +180,30 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta):
condition = kwargs.get("condition")
fallback_value = False
if None in [x_in, y_in, z_in, threshold, condition]:
logger.error(f"{self.__class__.__name__}: Missing required arguments.")
# Check required arguments
required_args = {
"x": x_in, "y": y_in, "z": z_in, "threshold": threshold, "condition": condition
}
if not coordinate_utils.check_required_arguments(required_args, self.__class__.__name__):
return (fallback_value,)
# --- Input Type Handling & History Cleanup ---
is_list = isinstance(x_in, list)
if is_list != isinstance(y_in, list) or is_list != isinstance(z_in, list):
raise ValueError(f"{self.__class__.__name__}: Inputs x, y, z must be all floats or all lists.")
# Validate input types
is_list = coordinate_utils.validate_coordinate_input_types(x_in, y_in, z_in, self.__class__.__name__)
if is_list:
current_batch_size = len(x_in)
# Remove histories for indices no longer present
keys_to_remove = [k for k in self._histories if k >= current_batch_size]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Clean up histories
coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__)
if current_batch_size == 0:
logger.warning(f"{self.__class__.__name__}: Received empty lists for coordinates.")
# Handle empty lists
if coordinate_utils.handle_empty_coordinate_lists(current_batch_size, self.__class__.__name__):
return ([],)
# Process List Input
results = []
for i in range(current_batch_size):
xi, yi, zi = x_in[i], y_in[i], z_in[i]
if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]):
logger.warning(f"{self.__class__.__name__}: Invalid coordinate type at index {i}. Using fallback value.")
if not coordinate_utils.validate_coordinate_types_in_list([[xi], [yi], [zi]], 0, self.__class__.__name__):
results.append(fallback_value)
continue
@@ -233,16 +222,11 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta):
return (results,)
else:
# Single Float Input: Remove all histories except for index 0
keys_to_remove = [k for k in self._histories if k != 0]
if keys_to_remove:
# logger.debug(f"{self.__class__.__name__}: Switching to single input, removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del self._histories[k]
# Clean up single input histories
coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__)
# Process Single Float Input
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
logger.error(f"{self.__class__.__name__}: Invalid coordinate types for single float input.")
# Validate single coordinate types
if not coordinate_utils.validate_single_coordinate_types(x_in, y_in, z_in, self.__class__.__name__):
return (fallback_value,)
raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0)
@@ -296,22 +280,18 @@ class CoordinateProximityNode:
y2: Union[float, List[float]],
z2: Union[float, List[float]],
) -> Union[float, List[float]]: # Return type hint updated
# --- Input Type Handling ---
# Validate inputs
inputs = [x1, y1, z1, x2, y2, z2]
is_list_input = isinstance(x1, list)
is_empty, list_len = coordinate_utils.validate_proximity_inputs(inputs, self.__class__.__name__)
fallback_value = 0.0
if is_empty:
return ([],) # Return empty list for FLOAT type
if any(isinstance(inp, list) != is_list_input for inp in inputs):
raise ValueError("All coordinate inputs (x1..z2) must be of the same type (all floats or all lists).")
is_list_input = isinstance(x1, list)
if is_list_input:
# Process List Input
list_len = len(x1)
if any(len(inp) != list_len for inp in inputs if isinstance(inp, list)):
raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.")
if list_len == 0:
logger.warning("Input coordinate lists are empty. Returning empty distance list.")
return ([],) # Return empty list for FLOAT type
num_calcs = list_len
x1_list, y1_list, z1_list = x1, y1, z1
x2_list, y2_list, z2_list = x2, y2, z2
@@ -399,28 +379,22 @@ class MaskFromCoordinate:
shape: str,
invert: bool,
) -> Tuple[torch.Tensor]:
# --- Input Type Handling ---
# Validate coordinate inputs
is_empty, num_coords = coordinate_utils.validate_mask_coordinate_inputs(x, y, self.__class__.__name__)
if is_empty:
# Determine shape for empty tensor
if image_for_dimensions.dim() != 4:
logger.error("Input image must be BHWC format to determine empty mask size.")
return (torch.zeros((0, 1, 64, 64), dtype=torch.float32),)
_, height, width, _ = image_for_dimensions.shape
return (torch.zeros((0, 1, height, width), dtype=torch.float32, device=image_for_dimensions.device),)
is_x_list = isinstance(x, list)
is_y_list = isinstance(y, list)
if is_x_list != is_y_list:
raise ValueError("Inputs x and y must be both floats or both lists.")
if is_x_list:
# Input is List[float]
x_list = x
y_list = y
if len(x_list) != len(y_list):
raise ValueError(f"x_list (len {len(x_list)}) and y_list (len {len(y_list)}) must have the same length.")
if not x_list: # Handle empty list
logger.warning("Input coordinate lists are empty. Returning empty mask batch.")
# Determine shape for empty tensor
if image_for_dimensions.dim() != 4:
logger.error("Input image must be BHWC format to determine empty mask size.")
return (torch.zeros((0, 1, 64, 64), dtype=torch.float32),)
_, height, width, _ = image_for_dimensions.shape
return (torch.zeros((0, 1, height, width), dtype=torch.float32, device=image_for_dimensions.device),)
num_coords = len(x_list)
else:
# Input is float - wrap in lists for unified processing
x_list = [x]
@@ -442,10 +416,7 @@ class MaskFromCoordinate:
yi = y_list[i]
# Validate individual coordinates
if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)):
logger.warning(
f"Invalid coordinate type at index {i} (x={type(xi)}, y={type(yi)}). Skipping mask generation for this index."
)
if not coordinate_utils.validate_individual_mask_coordinates(xi, yi, i, self.__class__.__name__):
mask_tensor = torch.zeros((1, height, width), dtype=torch.float32, device=device)
output_masks.append(mask_tensor)
continue
src/coordinates/__init__.py
+3 -1
View File
@@ -8,5 +8,7 @@ with support for single values, batches, and point collections.
from .coordinate_system import CoordinateSystem, Point, PointList
from .drawing_engine import DrawingEngine
from . import coordinate_utils
from . import drawing_utils
__all__ = ["CoordinateSystem", "Point", "PointList", "DrawingEngine"]
__all__ = ["CoordinateSystem", "Point", "PointList", "DrawingEngine", "coordinate_utils", "drawing_utils"]
src/coordinates/coordinate_utils.py
+138
View File
@@ -0,0 +1,138 @@
"""
Utility functions for coordinate processing
"""
import logging
from typing import List, Union, Tuple, Any, Dict
logger = logging.getLogger(__name__)
def validate_coordinate_input_types(x_in, y_in, z_in, class_name: str):
is_list = isinstance(x_in, list)
if is_list != isinstance(y_in, list) or is_list != isinstance(z_in, list):
raise ValueError(f"{class_name}: Inputs x, y, z must be all floats or all lists.")
return is_list
def cleanup_coordinate_histories(histories: Dict, current_batch_size: int, class_name: str):
keys_to_remove = [k for k in histories if k >= current_batch_size]
if keys_to_remove:
# logger.debug(f"{class_name}: Removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del histories[k]
def cleanup_single_input_histories(histories: Dict, class_name: str):
keys_to_remove = [k for k in histories if k != 0]
if keys_to_remove:
# logger.debug(f"{class_name}: Switching to single input, removing histories for indices {keys_to_remove}")
for k in keys_to_remove:
del histories[k]
def validate_coordinate_types_in_list(coords_list: List, index: int, class_name: str) -> bool:
xi, yi, zi = coords_list[0][index], coords_list[1][index], coords_list[2][index]
if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]):
logger.warning(f"{class_name}: Invalid coordinate type at index {index}. Using fallback value.")
return False
return True
def validate_single_coordinate_types(x_in, y_in, z_in, class_name: str) -> bool:
if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]):
logger.error(f"{class_name}: Invalid coordinate types for single float input.")
return False
return True
def check_required_arguments(args_dict: Dict, class_name: str) -> bool:
if None in args_dict.values():
logger.error(f"{class_name}: Missing required arguments.")
return False
return True
def handle_empty_coordinate_lists(current_batch_size: int, class_name: str) -> bool:
if current_batch_size == 0:
logger.warning(f"{class_name}: Received empty lists for coordinates.")
return True
return False
def standardize_inputs_to_lists(*inputs):
return [inp if isinstance(inp, list) else [inp] for inp in inputs]
def parse_label_values(label_values: str) -> List[str]:
values_list = []
if label_values and label_values.strip():
values_list = [v.strip() for v in label_values.split(',')]
return values_list
def determine_batch_label(values_list: List[str], batch_index: int, label_prefix: str, batch_value_mode: str) -> str:
if values_list:
if batch_value_mode == "Index-based" and batch_index < len(values_list):
# Use batch index to select value
current_value = values_list[batch_index]
else:
# Default to first value
current_value = values_list[0]
current_label = f"{label_prefix}: {current_value}"
else:
current_label = label_prefix
return current_label
def get_coordinate_space_and_dimensions(is_normalized: bool, image):
from . import CoordinateSystem
space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL
dimensions = CoordinateSystem.get_dimensions_from_tensor(image)
return space, dimensions
def validate_proximity_inputs(inputs: List, class_name: str) -> Tuple[bool, int]:
is_list_input = isinstance(inputs[0], list)
fallback_value = 0.0
if any(isinstance(inp, list) != is_list_input for inp in inputs):
raise ValueError("All coordinate inputs (x1..z2) must be of the same type (all floats or all lists).")
if is_list_input:
list_len = len(inputs[0])
if any(len(inp) != list_len for inp in inputs if isinstance(inp, list)):
raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.")
if list_len == 0:
logger.warning("Input coordinate lists are empty. Returning empty distance list.")
return True, 0 # Return empty list indicator
return False, list_len
return False, 1
def validate_mask_coordinate_inputs(x, y, class_name: str) -> Tuple[bool, int]:
is_x_list = isinstance(x, list)
is_y_list = isinstance(y, list)
if is_x_list != is_y_list:
raise ValueError("Inputs x and y must be both floats or both lists.")
if is_x_list:
x_list = x
y_list = y
if len(x_list) != len(y_list):
raise ValueError(f"x_list (len {len(x_list)}) and y_list (len {len(y_list)}) must have the same length.")
if not x_list: # Handle empty list
logger.warning("Input coordinate lists are empty. Returning empty mask batch.")
return True, 0 # Empty list indicator
return True, len(x_list)
return False, 1
def validate_individual_mask_coordinates(xi, yi, index: int, class_name: str) -> bool:
if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)):
logger.warning(
f"Invalid coordinate type at index {index} (x={type(xi)}, y={type(yi)}). Skipping mask generation for this index."
)
return False
return True
src/coordinates/drawing_utils.py
+68
View File
@@ -0,0 +1,68 @@
"""
Drawing utility functions
"""
import logging
from typing import List, Union, Tuple
import torch
logger = logging.getLogger(__name__)
def setup_batch_drawing(image: torch.Tensor):
batch_size = image.shape[0]
output_batch = image.clone()
return batch_size, output_batch
def handle_batch_mapping_coordinates(
batch_index: int,
batch_mapping: str,
coord_lists: List[List],
batch_size: int
) -> Tuple[bool, List]:
x_list, y_list = coord_lists[0], coord_lists[1]
if batch_mapping == "one-to-one":
# One-to-one: each batch item gets one coordinate
if batch_index < len(x_list):
batch_coords = [x_list[batch_index], y_list[batch_index]]
if len(coord_lists) > 2: # Handle additional coordinates (x2, y2 for lines)
batch_coords.extend([coord_list[batch_index] for coord_list in coord_lists[2:]])
return False, batch_coords
else:
return True, [] # Skip if not enough coordinates
elif batch_mapping == "all-on-first" and batch_index > 0:
return True, [] # Skip all but the first item
else:
# Broadcast: all coordinates on all batches
return False, coord_lists
def process_batch_item_drawing(
output_batch: torch.Tensor,
batch_index: int,
drawing_function,
drawing_args: dict
) -> torch.Tensor:
# Create a single-item tensor for this batch
single_image = output_batch[batch_index:batch_index+1]
# Process this batch item
result = drawing_function(image=single_image, **drawing_args)
# Update the output batch
output_batch[batch_index:batch_index+1] = result
return output_batch
def setup_drawing_engine_and_coordinates(image: torch.Tensor, is_normalized: bool):
from . import CoordinateSystem
from .drawing_engine import DrawingEngine
space = CoordinateSystem.NORMALIZED if is_normalized else CoordinateSystem.PIXEL
dimensions = CoordinateSystem.get_dimensions_from_tensor(image)
drawing_engine = DrawingEngine()
return drawing_engine, space, dimensions