From 8666aab3801fced29d166ce88af976ddbcf24a6d Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Wed, 28 May 2025 09:30:43 -0400 Subject: [PATCH] move drawing to gpu --- .../coordinates/position_tracking.py | 18 +- src/coordinates/coordinate_system.py | 155 +++--- src/coordinates/coordinate_utils.py | 146 ++---- src/coordinates/drawing_engine.py | 496 +++++++++++------- src/coordinates/drawing_utils.py | 51 +- 5 files changed, 467 insertions(+), 399 deletions(-) diff --git a/node_wrappers/coordinates/position_tracking.py b/node_wrappers/coordinates/position_tracking.py index 1483ce0..6224521 100644 --- a/node_wrappers/coordinates/position_tracking.py +++ b/node_wrappers/coordinates/position_tracking.py @@ -80,7 +80,7 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): "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__): + if None in required_args.values(): return (fallback_value,) # Validate input types @@ -92,14 +92,14 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__) # Handle empty lists - if coordinate_utils.handle_empty_coordinate_lists(current_batch_size, self.__class__.__name__): + if current_batch_size == 0: 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 coordinate_utils.validate_coordinate_types_in_list([[xi], [yi], [zi]], 0, self.__class__.__name__): + if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]): results.append(fallback_value) continue @@ -118,7 +118,7 @@ class CoordinateDeltaControlNode(BaseCoordinateDelta): coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__) # Validate single coordinate types - if not coordinate_utils.validate_single_coordinate_types(x_in, y_in, z_in, self.__class__.__name__): + if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]): return (fallback_value,) raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0) @@ -184,7 +184,7 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): 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__): + if None in required_args.values(): return (fallback_value,) # Validate input types @@ -196,14 +196,14 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): coordinate_utils.cleanup_coordinate_histories(self._histories, current_batch_size, self.__class__.__name__) # Handle empty lists - if coordinate_utils.handle_empty_coordinate_lists(current_batch_size, self.__class__.__name__): + if current_batch_size == 0: 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 coordinate_utils.validate_coordinate_types_in_list([[xi], [yi], [zi]], 0, self.__class__.__name__): + if not all(isinstance(c, (float, int)) for c in [xi, yi, zi]): results.append(fallback_value) continue @@ -226,7 +226,7 @@ class CoordinateDeltaTriggerNode(BaseCoordinateDelta): coordinate_utils.cleanup_single_input_histories(self._histories, self.__class__.__name__) # Validate single coordinate types - if not coordinate_utils.validate_single_coordinate_types(x_in, y_in, z_in, self.__class__.__name__): + if not all(isinstance(c, (float, int)) for c in [x_in, y_in, z_in]): return (fallback_value,) raw_delta = self._get_position_delta(float(x_in), float(y_in), float(z_in), window_size, item_index=0) @@ -416,7 +416,7 @@ class MaskFromCoordinate: yi = y_list[i] # Validate individual coordinates - if not coordinate_utils.validate_individual_mask_coordinates(xi, yi, i, self.__class__.__name__): + if not isinstance(xi, (float, int)) or not isinstance(yi, (float, int)): mask_tensor = torch.zeros((1, height, width), dtype=torch.float32, device=device) output_masks.append(mask_tensor) continue diff --git a/src/coordinates/coordinate_system.py b/src/coordinates/coordinate_system.py index 9e539ff..3d9069d 100644 --- a/src/coordinates/coordinate_system.py +++ b/src/coordinates/coordinate_system.py @@ -20,6 +20,24 @@ class CoordinateSystem: NORMALIZED = "normalized" # 0.0-1.0 range PIXEL = "pixel" # Pixel coordinates + # Cache for dimension tensors to avoid repeated creation + _dimension_cache = {} + + @staticmethod + def _get_dimension_tensor(dimensions: Union[int, Tuple[int, ...]], device: torch.device) -> torch.Tensor: + """Get cached dimension tensor or create new one.""" + cache_key = (dimensions, device) + if cache_key not in CoordinateSystem._dimension_cache: + if isinstance(dimensions, tuple): + CoordinateSystem._dimension_cache[cache_key] = torch.tensor(dimensions, device=device, dtype=torch.float32) + else: + CoordinateSystem._dimension_cache[cache_key] = torch.tensor([dimensions], device=device, dtype=torch.float32) + return CoordinateSystem._dimension_cache[cache_key] + + @staticmethod + def clear_cache(): + CoordinateSystem._dimension_cache.clear() + @staticmethod def normalize( coords: Union[float, List[float], torch.Tensor], @@ -40,36 +58,39 @@ class CoordinateSystem: if input_space == CoordinateSystem.NORMALIZED: return coords - # Handle different dimension specifications - if isinstance(dimensions, tuple): - dim_size = dimensions[0] if len(dimensions) == 1 else dimensions - else: - dim_size = dimensions - - # Handle different input types - if isinstance(coords, list): - if isinstance(dim_size, tuple): - # Multiple dimensions - if len(coords) % len(dim_size) != 0: - logger.warning(f"Coordinates length {len(coords)} not divisible by dimensions length {len(dim_size)}") + # Handle different input types with optimized paths + if isinstance(coords, torch.Tensor): + # Optimized tensor path - avoid device transfers + if isinstance(dimensions, tuple): + if coords.shape[-1] != len(dimensions): + raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dimensions)}") + dim_tensor = CoordinateSystem._get_dimension_tensor(dimensions, coords.device) + return coords.float() / dim_tensor + else: + return coords.float() / float(dimensions) + + elif isinstance(coords, list): + if not coords: # Empty list optimization + return coords + + if isinstance(dimensions, tuple): + # Vectorized list processing + dim_list = list(dimensions) result = [] for i, c in enumerate(coords): - dim_idx = i % len(dim_size) - result.append(float(c) / dim_size[dim_idx]) + dim_idx = i % len(dim_list) + result.append(float(c) / dim_list[dim_idx]) return result else: - # Single dimension - return [float(c) / dim_size for c in coords] - elif isinstance(coords, torch.Tensor): - if isinstance(dim_size, tuple): - # Tensor with shape handling - if coords.shape[-1] != len(dim_size): - raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dim_size)}") - return coords.float() / torch.tensor(dim_size, device=coords.device) - else: - return coords.float() / dim_size + # Single dimension - vectorized division + dim_float = float(dimensions) + return [float(c) / dim_float for c in coords] else: - return float(coords) / dim_size + # Single value + if isinstance(dimensions, tuple): + return float(coords) / dimensions[0] + else: + return float(coords) / dimensions @staticmethod def denormalize( @@ -91,37 +112,42 @@ class CoordinateSystem: if output_space == CoordinateSystem.NORMALIZED: return coords - # Handle different dimension specifications - if isinstance(dimensions, tuple): - dim_size = dimensions[0] if len(dimensions) == 1 else dimensions - else: - dim_size = dimensions - - # Handle different input types - if isinstance(coords, list): - if isinstance(dim_size, tuple): - # Multiple dimensions - if len(coords) % len(dim_size) != 0: - logger.warning(f"Coordinates length {len(coords)} not divisible by dimensions length {len(dim_size)}") - result = [] - for i, c in enumerate(coords): - dim_idx = i % len(dim_size) - result.append(float(max(0, min(dim_size[dim_idx] - 1, c * dim_size[dim_idx])))) - return result - else: - # Single dimension - return [float(max(0, min(dim_size - 1, c * dim_size))) for c in coords] - elif isinstance(coords, torch.Tensor): - if isinstance(dim_size, tuple): - # Tensor with shape handling - if coords.shape[-1] != len(dim_size): - raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dim_size)}") - dim_tensor = torch.tensor(dim_size, device=coords.device) + # Handle different input types with optimized paths + if isinstance(coords, torch.Tensor): + # Optimized tensor path + if isinstance(dimensions, tuple): + if coords.shape[-1] != len(dimensions): + raise ValueError(f"Tensor last dimension {coords.shape[-1]} doesn't match dimensions {len(dimensions)}") + dim_tensor = CoordinateSystem._get_dimension_tensor(dimensions, coords.device) return torch.clamp(coords * dim_tensor, 0, dim_tensor - 1) else: - return torch.clamp(coords * dim_size, 0, dim_size - 1) + dim_float = float(dimensions) + return torch.clamp(coords * dim_float, 0, dim_float - 1) + + elif isinstance(coords, list): + if not coords: # Empty list optimization + return coords + + if isinstance(dimensions, tuple): + # Vectorized list processing with bounds checking + dim_list = list(dimensions) + result = [] + for i, c in enumerate(coords): + dim_idx = i % len(dim_list) + dim_val = dim_list[dim_idx] + result.append(float(max(0, min(dim_val - 1, c * dim_val)))) + return result + else: + # Single dimension - vectorized processing + dim_float = float(dimensions) + return [float(max(0, min(dim_float - 1, c * dim_float))) for c in coords] else: - return float(max(0, min(dim_size - 1, coords * dim_size))) + # Single value + if isinstance(dimensions, tuple): + dim_val = dimensions[0] + return float(max(0, min(dim_val - 1, coords * dim_val))) + else: + return float(max(0, min(dimensions - 1, coords * dimensions))) @staticmethod def convert( @@ -145,17 +171,22 @@ class CoordinateSystem: if from_space == to_space: return coords - # First normalize if needed - if from_space != CoordinateSystem.NORMALIZED: - normalized = CoordinateSystem.normalize(coords, dimensions, from_space) + # Direct conversion path when possible + if from_space == CoordinateSystem.PIXEL and to_space == CoordinateSystem.NORMALIZED: + return CoordinateSystem.normalize(coords, dimensions, from_space) + elif from_space == CoordinateSystem.NORMALIZED and to_space == CoordinateSystem.PIXEL: + return CoordinateSystem.denormalize(coords, dimensions, to_space) else: - normalized = coords + # Two-step conversion (less common) + if from_space != CoordinateSystem.NORMALIZED: + normalized = CoordinateSystem.normalize(coords, dimensions, from_space) + else: + normalized = coords - # Then denormalize to target space if needed - if to_space != CoordinateSystem.NORMALIZED: - return CoordinateSystem.denormalize(normalized, dimensions, to_space) - else: - return normalized + if to_space != CoordinateSystem.NORMALIZED: + return CoordinateSystem.denormalize(normalized, dimensions, to_space) + else: + return normalized @staticmethod def get_dimensions_from_tensor(tensor: torch.Tensor) -> Tuple[int, int]: diff --git a/src/coordinates/coordinate_utils.py b/src/coordinates/coordinate_utils.py index 39444e5..c855bb7 100644 --- a/src/coordinates/coordinate_utils.py +++ b/src/coordinates/coordinate_utils.py @@ -9,80 +9,50 @@ 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): + x_is_list = isinstance(x_in, list) + y_is_list = isinstance(y_in, list) + z_is_list = isinstance(z_in, list) + if x_is_list != y_is_list or x_is_list != z_is_list: raise ValueError(f"{class_name}: Inputs x, y, z must be all floats or all lists.") - return is_list + return x_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: + if not histories: + return + for k in list(histories.keys()): + if k >= current_batch_size: 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 + if len(histories) <= 1: + return + keep_zero = 0 in histories + histories.clear() + if keep_zero: + pass # Will be recreated as needed def standardize_inputs_to_lists(*inputs): - return [inp if isinstance(inp, list) else [inp] for inp in inputs] + if all(isinstance(inp, list) for inp in inputs): + return inputs + return tuple(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 + if not label_values or not label_values.strip(): + return [] + return [v.strip() for v in label_values.split(',') if v.strip()] 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 + if not values_list: + return label_prefix + + if batch_value_mode == "Index-based" and batch_index < len(values_list): + return f"{label_prefix}: {values_list[batch_index]}" + return f"{label_prefix}: {values_list[0]}" def get_coordinate_space_and_dimensions(is_normalized: bool, image): @@ -93,46 +63,40 @@ def get_coordinate_space_and_dimensions(is_normalized: bool, image): def validate_proximity_inputs(inputs: List, class_name: str) -> Tuple[bool, int]: - is_list_input = isinstance(inputs[0], list) - fallback_value = 0.0 + if not inputs: + return True, 0 + + first_input = inputs[0] + is_list_input = isinstance(first_input, list) - 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).") + for inp in inputs[1:]: + if isinstance(inp, list) != is_list_input: + 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 + first_len = len(first_input) + if first_len == 0: + return True, 0 + + for inp in inputs[1:]: + if isinstance(inp, list) and len(inp) != first_len: + raise ValueError("If inputs are lists, all coordinate lists (x1..z2) must have the same length.") + + return False, first_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: + x_is_list = isinstance(x, list) + + if x_is_list != isinstance(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 + if x_is_list: + x_len = len(x) + if x_len != len(y): + raise ValueError(f"x_list (len {x_len}) and y_list (len {len(y)}) must have the same length.") + if x_len == 0: + return True, 0 + return True, x_len + return False, 1 \ No newline at end of file diff --git a/src/coordinates/drawing_engine.py b/src/coordinates/drawing_engine.py index 69bfcbd..a3cfa88 100644 --- a/src/coordinates/drawing_engine.py +++ b/src/coordinates/drawing_engine.py @@ -1,5 +1,6 @@ import logging from typing import List, Tuple, Union +import math import cv2 import numpy as np @@ -16,16 +17,154 @@ class DrawingEngine: """ @staticmethod - def hex_to_bgr(hex_color: str) -> Tuple[int, int, int]: - """Converts hex color to BGR tuple with validation.""" + def hex_to_rgb_tensor(hex_color: str, device: torch.device) -> torch.Tensor: hex_color = hex_color.lstrip("#") if len(hex_color) != 6: - return (0, 0, 0) # Default black + return torch.tensor([1.0, 0.0, 0.0], device=device, dtype=torch.float32) 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 + return torch.tensor([rgb[0]/255.0, rgb[1]/255.0, rgb[2]/255.0], device=device, dtype=torch.float32) except ValueError: - return (0, 0, 0) + return torch.tensor([1.0, 0.0, 0.0], device=device, dtype=torch.float32) + + @staticmethod + def create_coordinate_grids(height: int, width: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: + y_coords = torch.arange(height, device=device, dtype=torch.float32).view(-1, 1).expand(height, width) + x_coords = torch.arange(width, device=device, dtype=torch.float32).view(1, -1).expand(height, width) + return y_coords, x_coords + + @staticmethod + def draw_circles_gpu(image: torch.Tensor, coords: List[Tuple[int, int]], radius: int, color: torch.Tensor) -> torch.Tensor: + if not coords: + return image + + batch_size, height, width, channels = image.shape + device = image.device + + # Create coordinate grids once + y_grid, x_grid = DrawingEngine.create_coordinate_grids(height, width, device) + + for px, py in coords: + # Calculate squared distances (avoid sqrt for performance) + dist_sq = (x_grid - px) ** 2 + (y_grid - py) ** 2 + mask = dist_sq <= radius ** 2 + + # Apply color to all batch items and channels + for b in range(batch_size): + for c in range(channels): + image[b, :, :, c][mask] = color[c] + + return image + + @staticmethod + def draw_lines_gpu(image: torch.Tensor, lines: List[Tuple[int, int, int, int]], thickness: int, color: torch.Tensor) -> torch.Tensor: + if not lines: + return image + + batch_size, height, width, channels = image.shape + device = image.device + + # Create coordinate grids once + y_grid, x_grid = DrawingEngine.create_coordinate_grids(height, width, device) + + half_thickness = thickness / 2.0 + + for px1, py1, px2, py2 in lines: + # Handle degenerate cases + if px1 == px2 and py1 == py2: + # Single point + mask = (x_grid - px1) ** 2 + (y_grid - py1) ** 2 <= half_thickness ** 2 + else: + # Line segment using distance to line formula + # Vector from point1 to point2 + dx = px2 - px1 + dy = py2 - py1 + line_length_sq = dx * dx + dy * dy + + if line_length_sq == 0: + continue + + # Vector from point1 to each pixel + px_diff = x_grid - px1 + py_diff = y_grid - py1 + + # Project onto line segment (0 <= t <= 1) + t = torch.clamp((px_diff * dx + py_diff * dy) / line_length_sq, 0, 1) + + # Closest point on line segment + closest_x = px1 + t * dx + closest_y = py1 + t * dy + + # Distance from pixel to closest point on line + dist_sq = (x_grid - closest_x) ** 2 + (y_grid - closest_y) ** 2 + mask = dist_sq <= half_thickness ** 2 + + # Apply color to all batch items and channels + for b in range(batch_size): + for c in range(channels): + image[b, :, :, c][mask] = color[c] + + return image + + @staticmethod + def draw_polygon_gpu(image: torch.Tensor, points: List[Tuple[int, int]], fill: bool, thickness: int, color: torch.Tensor) -> torch.Tensor: + if len(points) < 3: + return image + + batch_size, height, width, channels = image.shape + device = image.device + + if fill: + # Vectorized point-in-polygon test using GPU + y_grid, x_grid = DrawingEngine.create_coordinate_grids(height, width, device) + + # Convert points to tensors + points_tensor = torch.tensor(points, device=device, dtype=torch.float32) + x_coords = points_tensor[:, 0] + y_coords = points_tensor[:, 1] + + # Initialize mask + mask = torch.zeros(height, width, device=device, dtype=torch.bool) + + # Vectorized ray casting algorithm + n_points = len(points) + for i in range(n_points): + j = (i - 1) % n_points + + xi, yi = x_coords[i], y_coords[i] + xj, yj = x_coords[j], y_coords[j] + + # Check if ray crosses edge + cond1 = (yi > y_grid) != (yj > y_grid) + + # Calculate intersection x-coordinate + denom = yj - yi + # Avoid division by zero + valid_denom = torch.abs(denom) > 1e-10 + intersection_x = torch.where( + valid_denom, + (xj - xi) * (y_grid - yi) / denom + xi, + torch.full_like(x_grid, float('inf')) + ) + + cond2 = x_grid < intersection_x + + # Toggle mask where both conditions are true + edge_cross = cond1 & cond2 & valid_denom + mask = mask ^ edge_cross + + # Apply color to filled area + for b in range(batch_size): + for c in range(channels): + image[b, :, :, c][mask] = color[c] + else: + # Draw outline by connecting line segments + for i in range(len(points)): + p1 = points[i] + p2 = points[(i + 1) % len(points)] + image = DrawingEngine.draw_lines_gpu(image, [(*p1, *p2)], thickness, color) + + return image @staticmethod def prep_batch_image(image: torch.Tensor, batch_idx: int) -> Tuple[np.ndarray, torch.device, torch.dtype]: @@ -34,11 +173,8 @@ class DrawingEngine: 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 @@ -46,6 +182,17 @@ class DrawingEngine: """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) + @staticmethod + def hex_to_bgr(hex_color: str) -> Tuple[int, int, int]: + hex_color = hex_color.lstrip("#") + if len(hex_color) != 6: + return (0, 0, 0) + try: + rgb = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) + return (rgb[2], rgb[1], rgb[0]) + except ValueError: + return (0, 0, 0) + def draw_points( self, image: torch.Tensor, @@ -71,66 +218,52 @@ class DrawingEngine: Modified image tensor """ batch_size, height, width, _ = image.shape - output_batch = image.clone() # Avoid modifying input - bgr_color = self.hex_to_bgr(color) + device = image.device + color_tensor = self.hex_to_rgb_tensor(color, device) - # 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 + # Pre-filter and convert coordinates + valid_coords = [] + for xi, yi in zip(x_list, y_list): + if isinstance(xi, (float, int)) and isinstance(yi, (float, int)): if is_normalized: - px = int(CoordinateSystem.denormalize(xi, width, CoordinateSystem.PIXEL)) - py = int(CoordinateSystem.denormalize(yi, height, CoordinateSystem.PIXEL)) + px = int(max(0, min(width - 1, CoordinateSystem.denormalize(xi, width, CoordinateSystem.PIXEL)))) + py = int(max(0, min(height - 1, CoordinateSystem.denormalize(yi, height, CoordinateSystem.PIXEL)))) else: px = int(max(0, min(width - 1, xi))) py = int(max(0, min(height - 1, yi))) + valid_coords.append((px, py)) - # Draw circle - cv2.circle(img_np, (px, py), radius, bgr_color, -1) + if not valid_coords: + return image - # Convert back to tensor and update batch - output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) + # Process each batch item + for b in range(batch_size): + if batch_mapping == "all-on-first" and b > 0: + continue - return output_batch + if batch_mapping == "one-to-one" and points_len == batch_size: + if b < len(valid_coords): + coords_to_draw = [valid_coords[b]] + else: + continue + else: + coords_to_draw = valid_coords + + if coords_to_draw: + # Extract single batch item for processing + batch_image = image[b:b+1] + batch_image = self.draw_circles_gpu(batch_image, coords_to_draw, radius, color_tensor) + image[b:b+1] = batch_image + + return image def draw_lines( self, @@ -150,120 +283,109 @@ class DrawingEngine: 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) + device = image.device + color_tensor = self.hex_to_rgb_tensor(color, device) - # 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 + # Pre-filter and convert coordinates + valid_lines = [] + endpoint_coords = [] + for i in range(line_len): + _x1, _y1, _x2, _y2 = x1_list[i], y1_list[i], x2_list[i], y2_list[i] + if all(isinstance(c, (float, int)) for c in [_x1, _y1, _x2, _y2]): if is_normalized: - px1 = int(CoordinateSystem.denormalize(_x1, width, CoordinateSystem.PIXEL)) - py1 = int(CoordinateSystem.denormalize(_y1, height, CoordinateSystem.PIXEL)) - px2 = int(CoordinateSystem.denormalize(_x2, width, CoordinateSystem.PIXEL)) - py2 = int(CoordinateSystem.denormalize(_y2, height, CoordinateSystem.PIXEL)) + px1 = int(max(0, min(width - 1, CoordinateSystem.denormalize(_x1, width, CoordinateSystem.PIXEL)))) + py1 = int(max(0, min(height - 1, CoordinateSystem.denormalize(_y1, height, CoordinateSystem.PIXEL)))) + px2 = int(max(0, min(width - 1, CoordinateSystem.denormalize(_x2, width, CoordinateSystem.PIXEL)))) + py2 = int(max(0, min(height - 1, CoordinateSystem.denormalize(_y2, height, CoordinateSystem.PIXEL)))) 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 + valid_lines.append((px1, py1, px2, py2)) 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) + endpoint_coords.extend([(px1, py1), (px2, py2)]) + if not valid_lines: + return image + + # Process each batch item + for b in range(batch_size): + if batch_mapping == "all-on-first" and b > 0: + continue + + if batch_mapping == "one-to-one" and line_len == batch_size: + if b < len(valid_lines): + lines_to_draw = [valid_lines[b]] + endpoints_to_draw = [] + if draw_endpoints: + line_data = valid_lines[b] + endpoints_to_draw = [(line_data[0], line_data[1]), (line_data[2], line_data[3])] + else: + continue + else: + lines_to_draw = valid_lines + endpoints_to_draw = endpoint_coords + + if lines_to_draw: + # Extract single batch item for processing + batch_image = image[b:b+1] + + # Draw lines using GPU + batch_image = self.draw_lines_gpu(batch_image, lines_to_draw, thickness, color_tensor) + + # Draw endpoints using GPU + if endpoints_to_draw: + batch_image = self.draw_circles_gpu(batch_image, endpoints_to_draw, point_radius, color_tensor) + + # Handle text labels (still use CV2 for now) 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" - label_x, label_y = px2, py2 + # Convert to numpy only for text rendering + img_np, device_orig, dtype_orig = self.prep_batch_image(batch_image, 0) + bgr_color = self.hex_to_bgr(color) + + for line_data in lines_to_draw: + px1, py1, px2, py2 = line_data + + 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: + label_x, label_y = px2, py2 - # Get text size - text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) + text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) - # Draw background rectangle - cv2.rectangle( - img_np, - (label_x - 5, label_y - text_size[1] - 5), - (label_x + text_size[0] + 5, label_y + 5), - (50, 50, 50), - -1, - ) + cv2.rectangle( + img_np, + (label_x - 5, label_y - text_size[1] - 5), + (label_x + text_size[0] + 5, label_y + 5), + (50, 50, 50), + -1, + ) - # Draw text - cv2.putText(img_np, label_text, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 1) + cv2.putText(img_np, label_text, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), 1) + + batch_image[0] = self.tensor_from_numpy(img_np, device_orig, dtype_orig) + + image[b:b+1] = batch_image - # Convert back to tensor and update batch - output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) - - return output_batch + return image def draw_polygon( self, @@ -282,95 +404,65 @@ class DrawingEngine: label_position: str = "Center", font_scale: float = 0.5, ) -> torch.Tensor: - """ - Draw polygon on images efficiently. - - Args: - image: BHWC tensor - points_x/y: Lists of polygon vertex coordinates - thickness: Line thickness in pixels - color: Hex color - fill: Whether to fill the polygon - is_normalized: If True, coords are 0-1 normalized - batch_mapping: How to map polygons to batch - 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) + device = image.device + color_tensor = self.hex_to_rgb_tensor(color, device) - # Validate input lengths if len(points_x) != len(points_y): raise ValueError("X and Y coordinate lists must have the same length") - # Ensure lists have at least 3 points for a polygon if len(points_x) < 3: logger.warning("Polygon needs at least 3 points. Drawing nothing.") - return output_batch + return image - # Determine batch mapping if batch_mapping == "one-to-one" and batch_size > 1: logger.warning(f"One-to-one mapping not supported for polygons. Using broadcast.") batch_mapping = "broadcast" - # 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 + # Pre-filter and convert coordinates + pixel_points = [] + for i in range(len(points_x)): + x_coord = points_x[i] + y_coord = points_y[i] - # Single efficient conversion to numpy - img_np, device, dtype = self.prep_batch_image(output_batch, b) - - # Convert polygon points to pixel coordinates - pixel_points = [] - for i in range(len(points_x)): - x_coord = points_x[i] - y_coord = points_y[i] - - if not isinstance(x_coord, (float, int)) or not isinstance(y_coord, (float, int)): - continue - - # Convert to pixel coordinates if needed + if isinstance(x_coord, (float, int)) and isinstance(y_coord, (float, int)): if is_normalized: - px = int(CoordinateSystem.denormalize(x_coord, width, CoordinateSystem.PIXEL)) - py = int(CoordinateSystem.denormalize(y_coord, height, CoordinateSystem.PIXEL)) + px = int(max(0, min(width - 1, CoordinateSystem.denormalize(x_coord, width, CoordinateSystem.PIXEL)))) + py = int(max(0, min(height - 1, CoordinateSystem.denormalize(y_coord, height, CoordinateSystem.PIXEL)))) else: px = int(max(0, min(width - 1, x_coord))) py = int(max(0, min(height - 1, y_coord))) - pixel_points.append((px, py)) - # Draw vertices if requested - if draw_vertices: - cv2.circle(img_np, (px, py), vertex_radius, bgr_color, -1) + if len(pixel_points) < 3: + return image - # Skip if no valid points - if len(pixel_points) < 3: + # Process each batch item + for b in range(batch_size): + if batch_mapping == "all-on-first" and b > 0: continue - # Convert to numpy array for OpenCV - pts = np.array(pixel_points, dtype=np.int32) + # Extract single batch item for processing + batch_image = image[b:b+1] + + # Draw polygon using GPU + batch_image = self.draw_polygon_gpu(batch_image, pixel_points, fill, thickness, color_tensor) + + # Draw vertices using GPU if requested + if draw_vertices: + batch_image = self.draw_circles_gpu(batch_image, pixel_points, vertex_radius, color_tensor) - # Draw filled or outline polygon - if fill: - cv2.fillPoly(img_np, [pts], bgr_color) - else: - cv2.polylines(img_np, [pts], True, bgr_color, thickness) - - # Draw label if requested + # Handle text labels (still use CV2 for now) if draw_label and label_text: - # Calculate label position (center of polygon) + # Convert to numpy only for text rendering + img_np, device_orig, dtype_orig = self.prep_batch_image(batch_image, 0) + bgr_color = self.hex_to_bgr(color) + center_x = int(sum(p[0] for p in pixel_points) / len(pixel_points)) center_y = int(sum(p[1] for p in pixel_points) / len(pixel_points)) - # Get text size text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1) - # Draw background rectangle cv2.rectangle( img_np, (center_x - text_size[0] // 2 - 5, center_y - text_size[1] // 2 - 5), @@ -379,7 +471,6 @@ class DrawingEngine: -1, ) - # Draw text cv2.putText( img_np, label_text, @@ -389,8 +480,9 @@ class DrawingEngine: (255, 255, 255), 1, ) + + batch_image[0] = self.tensor_from_numpy(img_np, device_orig, dtype_orig) - # Convert back to tensor and update batch - output_batch[b] = self.tensor_from_numpy(img_np, device, dtype) + image[b:b+1] = batch_image - return output_batch + return image diff --git a/src/coordinates/drawing_utils.py b/src/coordinates/drawing_utils.py index 3ca30ae..2bbffce 100644 --- a/src/coordinates/drawing_utils.py +++ b/src/coordinates/drawing_utils.py @@ -10,9 +10,7 @@ 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 + return image.shape[0], image.clone() def handle_batch_mapping_coordinates( @@ -21,40 +19,24 @@ def handle_batch_mapping_coordinates( coord_lists: List[List], batch_size: int ) -> Tuple[bool, List]: - x_list, y_list = coord_lists[0], coord_lists[1] + if batch_mapping == "all-on-first" and batch_index > 0: + return True, [] 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 + x_list, y_list = coord_lists[0], coord_lists[1] + if batch_index >= len(x_list): + return True, [] + + # Build coordinates list efficiently + batch_coords = [x_list[batch_index], y_list[batch_index]] + if len(coord_lists) > 2: + batch_coords.extend(coord_list[batch_index] for coord_list in coord_lists[2:]) + return False, batch_coords + + # Broadcast case + 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): @@ -63,6 +45,5 @@ def setup_drawing_engine_and_coordinates(image: torch.Tensor, is_normalized: boo 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 + return DrawingEngine(), space, dimensions \ No newline at end of file