move drawing to gpu

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

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]:

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 
+    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 

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

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 
+    return DrawingEngine(), space, dimensions