489 lines
20 KiB
Python
489 lines
20 KiB
Python
import logging
|
|
from typing import List, Tuple, Union
|
|
import math
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .coordinate_system import CoordinateSystem
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DrawingEngine:
|
|
"""High-performance drawing engine for real-time applications.
|
|
Optimized for minimal tensor-numpy conversions and efficient batch operations.
|
|
"""
|
|
|
|
@staticmethod
|
|
def hex_to_rgb_tensor(hex_color: str, device: torch.device) -> torch.Tensor:
|
|
hex_color = hex_color.lstrip("#")
|
|
if len(hex_color) != 6:
|
|
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 torch.tensor([rgb[0]/255.0, rgb[1]/255.0, rgb[2]/255.0], device=device, dtype=torch.float32)
|
|
except ValueError:
|
|
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]:
|
|
"""Fast preparation of single image from batch for drawing - minimal conversion."""
|
|
if image.dim() != 4:
|
|
raise ValueError("Input image must be BHWC format")
|
|
|
|
device, dtype = image.device, image.dtype
|
|
img_tensor = image[batch_idx].to(torch.float32)
|
|
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
return np.ascontiguousarray(img_np), device, dtype
|
|
|
|
@staticmethod
|
|
def tensor_from_numpy(img_np: np.ndarray, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
|
"""Fast conversion from numpy back to tensor with proper device/dtype."""
|
|
return torch.from_numpy(img_np.astype(np.float32) / 255.0).to(device=device, dtype=dtype)
|
|
|
|
@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,
|
|
points_x: Union[float, List[float]],
|
|
points_y: Union[float, List[float]],
|
|
radius: int = 5,
|
|
color: str = "#FF0000",
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
) -> torch.Tensor:
|
|
"""
|
|
Draw points on images efficiently.
|
|
|
|
Args:
|
|
image: BHWC tensor
|
|
points_x/y: Coordinate values (normalized 0-1 or pixel)
|
|
radius: Point radius in pixels
|
|
color: Hex color
|
|
is_normalized: If True, coords are 0-1 normalized
|
|
batch_mapping: How to map points to batch ("broadcast", "one-to-one", "all-on-first")
|
|
|
|
Returns:
|
|
Modified image tensor
|
|
"""
|
|
batch_size, height, width, _ = image.shape
|
|
device = image.device
|
|
color_tensor = self.hex_to_rgb_tensor(color, device)
|
|
|
|
x_list = points_x if isinstance(points_x, list) else [points_x]
|
|
y_list = points_y if isinstance(points_y, list) else [points_y]
|
|
|
|
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"
|
|
|
|
# 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(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))
|
|
|
|
if not valid_coords:
|
|
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 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,
|
|
image: torch.Tensor,
|
|
x1: Union[float, List[float]],
|
|
y1: Union[float, List[float]],
|
|
x2: Union[float, List[float]],
|
|
y2: Union[float, List[float]],
|
|
thickness: int = 2,
|
|
color: str = "#00FF00",
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
draw_endpoints: bool = False,
|
|
point_radius: int = 3,
|
|
draw_label: bool = False,
|
|
label_text: str = "",
|
|
label_position: str = "Midpoint",
|
|
font_scale: float = 0.5,
|
|
) -> torch.Tensor:
|
|
batch_size, height, width, _ = image.shape
|
|
device = image.device
|
|
color_tensor = self.hex_to_rgb_tensor(color, device)
|
|
|
|
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]
|
|
|
|
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")
|
|
|
|
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"
|
|
|
|
# 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(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)))
|
|
valid_lines.append((px1, py1, px2, py2))
|
|
if draw_endpoints:
|
|
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:
|
|
# 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
|
|
|
|
text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 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,
|
|
)
|
|
|
|
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
|
|
|
|
return image
|
|
|
|
def draw_polygon(
|
|
self,
|
|
image: torch.Tensor,
|
|
points_x: List[float],
|
|
points_y: List[float],
|
|
color: str = "#0000FF",
|
|
thickness: int = 2,
|
|
fill: bool = False,
|
|
is_normalized: bool = True,
|
|
batch_mapping: str = "broadcast",
|
|
draw_vertices: bool = False,
|
|
vertex_radius: int = 3,
|
|
draw_label: bool = False,
|
|
label_text: str = "",
|
|
label_position: str = "Center",
|
|
font_scale: float = 0.5,
|
|
) -> torch.Tensor:
|
|
batch_size, height, width, _ = image.shape
|
|
device = image.device
|
|
color_tensor = self.hex_to_rgb_tensor(color, device)
|
|
|
|
if len(points_x) != len(points_y):
|
|
raise ValueError("X and Y coordinate lists must have the same length")
|
|
|
|
if len(points_x) < 3:
|
|
logger.warning("Polygon needs at least 3 points. Drawing nothing.")
|
|
return image
|
|
|
|
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"
|
|
|
|
# Pre-filter and convert coordinates
|
|
pixel_points = []
|
|
for i in range(len(points_x)):
|
|
x_coord = points_x[i]
|
|
y_coord = points_y[i]
|
|
|
|
if isinstance(x_coord, (float, int)) and isinstance(y_coord, (float, int)):
|
|
if is_normalized:
|
|
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))
|
|
|
|
if len(pixel_points) < 3:
|
|
return image
|
|
|
|
# Process each batch item
|
|
for b in range(batch_size):
|
|
if batch_mapping == "all-on-first" and b > 0:
|
|
continue
|
|
|
|
# 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)
|
|
|
|
# Handle text labels (still use CV2 for now)
|
|
if draw_label and label_text:
|
|
# 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))
|
|
|
|
text_size, baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1)
|
|
|
|
cv2.rectangle(
|
|
img_np,
|
|
(center_x - text_size[0] // 2 - 5, center_y - text_size[1] // 2 - 5),
|
|
(center_x + text_size[0] // 2 + 5, center_y + text_size[1] // 2 + 5),
|
|
(50, 50, 50),
|
|
-1,
|
|
)
|
|
|
|
cv2.putText(
|
|
img_np,
|
|
label_text,
|
|
(center_x - text_size[0] // 2, center_y + text_size[1] // 2),
|
|
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
|
|
|
|
return image
|