ykcheungab
126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def _prepare_draw_mask(valid_mask, shape):
|
|
"""Resize a boolean mask to (height, width) for arrow culling."""
|
|
if valid_mask is None:
|
|
return None
|
|
h, w = shape[:2]
|
|
mask = np.asarray(valid_mask, dtype=bool)
|
|
if mask.shape != (h, w):
|
|
mask_u8 = mask.astype(np.uint8)
|
|
if mask_u8.ndim > 2:
|
|
mask_u8 = mask_u8[:, :, 0]
|
|
mask = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
|
|
return mask
|
|
|
|
|
|
def draw_arrows(
|
|
frame,
|
|
optical_flow,
|
|
fps=True,
|
|
threshold=5,
|
|
grid_spacing=40,
|
|
arrow_scale=1.0,
|
|
margin=0.8,
|
|
colormap=cv2.COLORMAP_JET,
|
|
below_threshold_color=(0, 0, 255),
|
|
edge_margin=15,
|
|
valid_mask=None,
|
|
):
|
|
"""
|
|
Draw colored optical flow arrows on the image by magnitude (JET colormap).
|
|
Frame and optical_flow must have the same height/width; arrows are drawn in pixel coordinates.
|
|
|
|
Parameters
|
|
----------
|
|
valid_mask : array-like, optional
|
|
Boolean (H, W) mask. Grid points outside the mask are skipped (no arrows or dots).
|
|
"""
|
|
image = frame.copy()
|
|
if len(image.shape) == 2 or image.shape[2] == 1:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
|
|
|
h, w = image.shape[:2]
|
|
mask = _prepare_draw_mask(valid_mask, (h, w))
|
|
|
|
# Arrow grid in pixel coordinates (assume frame and flow same size)
|
|
rr = np.arange(grid_spacing // 2, h, grid_spacing)
|
|
cc = np.arange(grid_spacing // 2, w, grid_spacing)
|
|
rr, cc = np.meshgrid(rr, cc, indexing="ij")
|
|
rr, cc = rr.reshape(-1), cc.reshape(-1)
|
|
|
|
# Displacement at each grid point (flow stored as half-displacement)
|
|
disp = 2.0 * optical_flow[rr, cc, :]
|
|
mag = np.linalg.norm(disp, axis=1)
|
|
dx = disp[:, 0] / arrow_scale
|
|
dy = disp[:, 1] / arrow_scale
|
|
length = np.sqrt(dx * dx + dy * dy)
|
|
# Cap arrow length at grid_spacing
|
|
scale = np.where(length > grid_spacing * margin, grid_spacing / np.maximum(length, 1e-9), 1.0) * margin
|
|
dx, dy = dx * scale, dy * scale
|
|
end_x = (cc + dx).astype(np.int32)
|
|
end_y = (rr + dy).astype(np.int32)
|
|
|
|
valid = (
|
|
(cc >= edge_margin)
|
|
& (cc < w - edge_margin)
|
|
& (rr >= edge_margin)
|
|
& (rr < h - edge_margin)
|
|
& (end_x >= 0)
|
|
& (end_x < w)
|
|
& (end_y >= 0)
|
|
& (end_y < h)
|
|
)
|
|
if mask is not None:
|
|
valid &= mask[rr, cc]
|
|
valid &= mask[end_y, end_x]
|
|
|
|
mag_range = mag.max() - mag.min()
|
|
norm_mag = (
|
|
((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
|
|
)
|
|
norm_mag[mag < threshold] = 0
|
|
colors = cv2.applyColorMap(norm_mag.reshape(-1, 1), colormap)
|
|
|
|
for i in np.where(valid)[0]:
|
|
pt1, pt2 = (int(cc[i]), int(rr[i])), (int(end_x[i]), int(end_y[i]))
|
|
color = tuple(int(c) for c in colors[i, 0])
|
|
if mag[i] < threshold:
|
|
cv2.circle(image, pt1, radius=1, color=below_threshold_color, thickness=-1)
|
|
else:
|
|
thickness = 1
|
|
cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
|
|
return image
|
|
|
|
|
|
def draw_contact(image, contour, centroid, point_color=(0, 255, 0), region_color=(0, 255, 255)):
|
|
"""
|
|
在图像上绘制扩张焦点和扩张区域。
|
|
|
|
Args:
|
|
image: 背景图像
|
|
contour: 轮廓
|
|
centroid: 触碰点
|
|
point_color: 触碰点的颜色 (BGR)
|
|
region_color: 轮廓的颜色 (BGR)
|
|
|
|
Returns:
|
|
np.ndarray: 绘制了触碰点和轮廓的图像
|
|
"""
|
|
image = image.copy()
|
|
|
|
# 绘制扩张区域轮廓
|
|
if contour is not None:
|
|
cv2.drawContours(image, [contour], -1, region_color, 2)
|
|
|
|
# 绘制触碰点
|
|
if centroid is not None:
|
|
cx, cy = centroid
|
|
cv2.circle(image, (cx, cy), 8, point_color, -1)
|
|
cv2.circle(image, (cx, cy), 10, point_color, 2)
|
|
|
|
return image
|
|
|