2026-05-18-17-40-02 构建导丝分割Web系统
This commit is contained in:
232
backend/segmentation.py
Normal file
232
backend/segmentation.py
Normal file
@@ -0,0 +1,232 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from skimage.filters import frangi, threshold_otsu
|
||||
from skimage.morphology import remove_small_objects, skeletonize
|
||||
|
||||
|
||||
METHOD_DESCRIPTIONS = {
|
||||
"hessian_ridge": {
|
||||
"label": "Hessian / Frangi 细线增强",
|
||||
"description": "多尺度 Hessian 管状结构响应，适合低对比细导丝候选提取。",
|
||||
"uses_temporal": False,
|
||||
},
|
||||
"edge_morphology": {
|
||||
"label": "边缘 + 形态学",
|
||||
"description": "CLAHE、黑帽增强、Canny 边缘与线性形态学连接。",
|
||||
"uses_temporal": False,
|
||||
},
|
||||
"temporal_difference": {
|
||||
"label": "视频时序差分",
|
||||
"description": "利用相邻帧运动候选抑制静态骨骼和背景结构。",
|
||||
"uses_temporal": True,
|
||||
},
|
||||
"fusion": {
|
||||
"label": "融合模式",
|
||||
"description": "融合 Hessian、边缘形态学和时序差分，作为默认稳健输出。",
|
||||
"uses_temporal": True,
|
||||
},
|
||||
"compare": {
|
||||
"label": "多方法对比",
|
||||
"description": "对同一帧同时运行多种方法，便于调参和方案比较。",
|
||||
"uses_temporal": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SegmentationOutput:
|
||||
method: str
|
||||
mask: np.ndarray
|
||||
overlay: np.ndarray
|
||||
metrics: dict[str, float | int]
|
||||
|
||||
|
||||
def normalize01(image: np.ndarray) -> np.ndarray:
|
||||
image = image.astype(np.float32)
|
||||
low = float(np.percentile(image, 1))
|
||||
high = float(np.percentile(image, 99))
|
||||
if high <= low:
|
||||
return np.zeros_like(image, dtype=np.float32)
|
||||
return np.clip((image - low) / (high - low), 0.0, 1.0)
|
||||
|
||||
|
||||
def to_gray(frame: np.ndarray) -> np.ndarray:
|
||||
if frame.ndim == 2:
|
||||
return frame
|
||||
return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
|
||||
def clahe_gray(frame: np.ndarray) -> np.ndarray:
|
||||
gray = to_gray(frame)
|
||||
clahe = cv2.createCLAHE(clipLimit=2.2, tileGridSize=(8, 8))
|
||||
return clahe.apply(gray)
|
||||
|
||||
|
||||
def _adaptive_cutoff(response: np.ndarray, sensitivity: float) -> float:
|
||||
response = normalize01(response)
|
||||
nonzero = response[response > 0]
|
||||
if nonzero.size < 16:
|
||||
return 1.0
|
||||
sensitivity = float(np.clip(sensitivity, 0.05, 0.95))
|
||||
percentile = 99.2 - sensitivity * 22.0
|
||||
percentile_cut = float(np.percentile(nonzero, percentile))
|
||||
try:
|
||||
otsu_cut = float(threshold_otsu(nonzero))
|
||||
except ValueError:
|
||||
otsu_cut = percentile_cut
|
||||
return max(min(percentile_cut, 0.98), otsu_cut * 0.72)
|
||||
|
||||
|
||||
def clean_mask(mask: np.ndarray, min_area: int = 12) -> np.ndarray:
|
||||
binary = mask.astype(bool)
|
||||
binary = remove_small_objects(binary, max_size=max(1, int(min_area)))
|
||||
cleaned = binary.astype(np.uint8) * 255
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel, iterations=1)
|
||||
return cleaned
|
||||
|
||||
|
||||
def hessian_ridge_mask(frame: np.ndarray, sensitivity: float = 0.56) -> np.ndarray:
|
||||
enhanced = clahe_gray(frame)
|
||||
inverted = 255 - enhanced
|
||||
normalized = normalize01(inverted)
|
||||
response = frangi(
|
||||
normalized,
|
||||
sigmas=(0.7, 1.1, 1.7, 2.3),
|
||||
alpha=0.55,
|
||||
beta=0.55,
|
||||
gamma=12,
|
||||
black_ridges=False,
|
||||
)
|
||||
response = normalize01(response)
|
||||
cutoff = _adaptive_cutoff(response, sensitivity)
|
||||
mask = response >= cutoff
|
||||
return clean_mask(mask, min_area=10)
|
||||
|
||||
|
||||
def edge_morphology_mask(frame: np.ndarray, sensitivity: float = 0.56) -> np.ndarray:
|
||||
enhanced = clahe_gray(frame)
|
||||
blur = cv2.GaussianBlur(enhanced, (3, 3), 0)
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
|
||||
blackhat = cv2.morphologyEx(blur, cv2.MORPH_BLACKHAT, kernel)
|
||||
dark_line = normalize01(blackhat)
|
||||
cutoff = _adaptive_cutoff(dark_line, min(0.95, sensitivity + 0.1))
|
||||
candidate = (dark_line >= cutoff).astype(np.uint8) * 255
|
||||
low = int(20 + (1.0 - sensitivity) * 65)
|
||||
high = int(70 + (1.0 - sensitivity) * 120)
|
||||
edges = cv2.Canny(blur, low, high)
|
||||
candidate = cv2.dilate(candidate, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
|
||||
edges = cv2.bitwise_and(edges, candidate)
|
||||
line_h = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 1))
|
||||
line_v = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 5))
|
||||
connected = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, line_h, iterations=1)
|
||||
connected = cv2.morphologyEx(connected, cv2.MORPH_CLOSE, line_v, iterations=1)
|
||||
connected = cv2.bitwise_or(connected, cv2.erode(candidate, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))))
|
||||
return clean_mask(connected, min_area=8)
|
||||
|
||||
|
||||
def temporal_difference_mask(
|
||||
frame: np.ndarray,
|
||||
previous_frame: np.ndarray | None,
|
||||
sensitivity: float = 0.56,
|
||||
) -> np.ndarray:
|
||||
ridge = hessian_ridge_mask(frame, sensitivity=sensitivity)
|
||||
if previous_frame is None:
|
||||
return ridge
|
||||
current = cv2.GaussianBlur(clahe_gray(frame), (5, 5), 0)
|
||||
previous = cv2.GaussianBlur(clahe_gray(previous_frame), (5, 5), 0)
|
||||
diff = cv2.absdiff(current, previous)
|
||||
diff = normalize01(diff)
|
||||
cutoff = _adaptive_cutoff(diff, min(0.92, sensitivity + 0.12))
|
||||
moving = (diff >= cutoff).astype(np.uint8) * 255
|
||||
moving = cv2.dilate(moving, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
|
||||
blended = cv2.bitwise_or(cv2.bitwise_and(ridge, moving), cv2.bitwise_and(ridge, cv2.dilate(moving, None)))
|
||||
if int(np.count_nonzero(blended)) < 8:
|
||||
blended = cv2.bitwise_or(ridge, moving)
|
||||
return clean_mask(blended, min_area=8)
|
||||
|
||||
|
||||
def fusion_mask(
|
||||
frame: np.ndarray,
|
||||
previous_frame: np.ndarray | None = None,
|
||||
sensitivity: float = 0.56,
|
||||
) -> np.ndarray:
|
||||
ridge = hessian_ridge_mask(frame, sensitivity=sensitivity)
|
||||
edge = edge_morphology_mask(frame, sensitivity=sensitivity)
|
||||
temporal = temporal_difference_mask(frame, previous_frame, sensitivity=sensitivity)
|
||||
votes = (
|
||||
(ridge > 0).astype(np.uint8)
|
||||
+ (edge > 0).astype(np.uint8)
|
||||
+ (temporal > 0).astype(np.uint8)
|
||||
)
|
||||
fused = votes >= 2
|
||||
if int(np.count_nonzero(fused)) < 8:
|
||||
fused = votes >= 1
|
||||
return clean_mask(fused, min_area=10)
|
||||
|
||||
|
||||
def overlay_mask(frame: np.ndarray, mask: np.ndarray, color: tuple[int, int, int] = (0, 220, 255)) -> np.ndarray:
|
||||
if frame.ndim == 2:
|
||||
base = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
|
||||
else:
|
||||
base = frame.copy()
|
||||
color_layer = np.zeros_like(base)
|
||||
color_layer[mask > 0] = color
|
||||
return cv2.addWeighted(base, 0.78, color_layer, 0.72, 0)
|
||||
|
||||
|
||||
def mask_metrics(mask: np.ndarray) -> dict[str, float | int]:
|
||||
binary = mask > 0
|
||||
coverage = float(np.count_nonzero(binary) / binary.size) if binary.size else 0.0
|
||||
skeleton = skeletonize(binary)
|
||||
component_count, _ = cv2.connectedComponents(binary.astype(np.uint8))
|
||||
return {
|
||||
"coverage": round(coverage, 6),
|
||||
"mask_pixels": int(np.count_nonzero(binary)),
|
||||
"skeleton_length": int(np.count_nonzero(skeleton)),
|
||||
"components": max(0, int(component_count) - 1),
|
||||
}
|
||||
|
||||
|
||||
def segment_frame(
|
||||
frame: np.ndarray,
|
||||
method: str = "fusion",
|
||||
previous_frame: np.ndarray | None = None,
|
||||
sensitivity: float = 0.56,
|
||||
) -> SegmentationOutput:
|
||||
method_map: dict[str, Callable[..., np.ndarray]] = {
|
||||
"hessian_ridge": hessian_ridge_mask,
|
||||
"edge_morphology": edge_morphology_mask,
|
||||
"temporal_difference": temporal_difference_mask,
|
||||
"fusion": fusion_mask,
|
||||
}
|
||||
if method not in method_map:
|
||||
raise ValueError(f"Unknown segmentation method: {method}")
|
||||
if method in {"temporal_difference", "fusion"}:
|
||||
mask = method_map[method](frame, previous_frame, sensitivity)
|
||||
else:
|
||||
mask = method_map[method](frame, sensitivity)
|
||||
return SegmentationOutput(
|
||||
method=method,
|
||||
mask=mask,
|
||||
overlay=overlay_mask(frame, mask),
|
||||
metrics=mask_metrics(mask),
|
||||
)
|
||||
|
||||
|
||||
def compare_frame(
|
||||
frame: np.ndarray,
|
||||
previous_frame: np.ndarray | None = None,
|
||||
sensitivity: float = 0.56,
|
||||
) -> list[SegmentationOutput]:
|
||||
return [
|
||||
segment_frame(frame, "hessian_ridge", previous_frame, sensitivity),
|
||||
segment_frame(frame, "edge_morphology", previous_frame, sensitivity),
|
||||
segment_frame(frame, "temporal_difference", previous_frame, sensitivity),
|
||||
segment_frame(frame, "fusion", previous_frame, sensitivity),
|
||||
]
|
||||
Reference in New Issue
Block a user