From 3a8a798a7fac1ce5ea9bee8881dc3879fec33dfc Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Tue, 27 May 2025 17:14:39 -0400 Subject: [PATCH] code reduction --- node_wrappers/coordinates/drawing_nodes.py | 127 ++++++---------- .../coordinates/position_tracking.py | 133 +++++++---------- src/coordinates/__init__.py | 4 +- src/coordinates/coordinate_utils.py | 138 ++++++++++++++++++ src/coordinates/drawing_utils.py | 68 +++++++++ 5 files changed, 306 insertions(+), 164 deletions(-) create mode 100644 src/coordinates/coordinate_utils.py create mode 100644 src/coordinates/drawing_utils.py diff --git a/node_wrappers/coordinates/drawing_nodes.py b/node_wrappers/coordinates/drawing_nodes.py index 92f7da8..67dd0ee 100644 --- a/node_wrappers/coordinates/drawing_nodes.py +++ b/node_wrappers/coordinates/drawing_nodes.py @@ -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) diff --git a/node_wrappers/coordinates/position_tracking.py b/node_wrappers/coordinates/position_tracking.py index fc9723e..1483ce0 100644 --- a/node_wrappers/coordinates/position_tracking.py +++ b/node_wrappers/coordinates/position_tracking.py @@ -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 diff --git a/src/coordinates/__init__.py b/src/coordinates/__init__.py index d23a798..90dddf0 100644 --- a/src/coordinates/__init__.py +++ b/src/coordinates/__init__.py @@ -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"] diff --git a/src/coordinates/coordinate_utils.py b/src/coordinates/coordinate_utils.py new file mode 100644 index 0000000..39444e5 --- /dev/null +++ b/src/coordinates/coordinate_utils.py @@ -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 \ No newline at end of file diff --git a/src/coordinates/drawing_utils.py b/src/coordinates/drawing_utils.py new file mode 100644 index 0000000..3ca30ae --- /dev/null +++ b/src/coordinates/drawing_utils.py @@ -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 \ No newline at end of file