c20d6e55b6
- engine/pipeline/preprocessing.py - engine/tests/test_preprocessing.py GSD-Task: S01/T02
130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
"""OpenCV preprocessing pipeline for raster-to-vector conversion."""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def decode_image(raw_bytes: bytes) -> np.ndarray:
|
|
"""Decode raw image bytes into a BGR numpy array."""
|
|
buf = np.frombuffer(raw_bytes, dtype=np.uint8)
|
|
img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
raise ValueError("Failed to decode image from provided bytes")
|
|
return img
|
|
|
|
|
|
def to_grayscale(img: np.ndarray) -> np.ndarray:
|
|
"""Convert BGR image to single-channel grayscale."""
|
|
if len(img.shape) == 2:
|
|
return img
|
|
return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
|
|
def denoise(
|
|
img: np.ndarray,
|
|
d: int = 9,
|
|
sigma_color: float = 75.0,
|
|
sigma_space: float = 75.0,
|
|
) -> np.ndarray:
|
|
"""Apply bilateral filter for edge-preserving denoising."""
|
|
return cv2.bilateralFilter(img, d, sigma_color, sigma_space)
|
|
|
|
|
|
def enhance_contrast(
|
|
img: np.ndarray,
|
|
clip_limit: float = 2.0,
|
|
tile_grid_size: tuple[int, int] = (8, 8),
|
|
) -> np.ndarray:
|
|
"""Apply CLAHE contrast enhancement."""
|
|
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
|
|
return clahe.apply(img)
|
|
|
|
|
|
def threshold(
|
|
img: np.ndarray,
|
|
manual_thresh: int | None = None,
|
|
) -> np.ndarray:
|
|
"""Apply thresholding — Otsu auto by default, manual override if provided."""
|
|
if manual_thresh is not None:
|
|
_, result = cv2.threshold(img, manual_thresh, 255, cv2.THRESH_BINARY)
|
|
else:
|
|
_, result = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
return result
|
|
|
|
|
|
def edge_detect(
|
|
img: np.ndarray,
|
|
low: int = 50,
|
|
high: int = 150,
|
|
) -> np.ndarray:
|
|
"""Apply Canny edge detection."""
|
|
return cv2.Canny(img, low, high)
|
|
|
|
|
|
def morphological_ops(
|
|
img: np.ndarray,
|
|
kernel_size: int = 3,
|
|
dilate_iterations: int = 1,
|
|
erode_iterations: int = 1,
|
|
) -> np.ndarray:
|
|
"""Apply dilation then erosion (closing-style) to clean up binary image."""
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
|
|
result = cv2.dilate(img, kernel, iterations=dilate_iterations)
|
|
result = cv2.erode(result, kernel, iterations=erode_iterations)
|
|
return result
|
|
|
|
|
|
def preprocess(
|
|
raw_bytes: bytes,
|
|
params: dict | None = None,
|
|
) -> np.ndarray:
|
|
"""Run the full preprocessing pipeline on raw image bytes.
|
|
|
|
Stages: decode → grayscale → denoise → contrast → threshold → morphological ops.
|
|
Edge detection is optional (enabled via params["edge_detect"] = True).
|
|
|
|
All stage parameters are tunable via the params dict. Keys:
|
|
denoise_d, denoise_sigma_color, denoise_sigma_space,
|
|
clahe_clip_limit, clahe_tile_grid_size,
|
|
threshold_manual,
|
|
edge_detect (bool), edge_low, edge_high,
|
|
morph_kernel_size, morph_dilate_iterations, morph_erode_iterations
|
|
"""
|
|
p = params or {}
|
|
|
|
img = decode_image(raw_bytes)
|
|
img = to_grayscale(img)
|
|
|
|
img = denoise(
|
|
img,
|
|
d=p.get("denoise_d", 9),
|
|
sigma_color=p.get("denoise_sigma_color", 75.0),
|
|
sigma_space=p.get("denoise_sigma_space", 75.0),
|
|
)
|
|
|
|
img = enhance_contrast(
|
|
img,
|
|
clip_limit=p.get("clahe_clip_limit", 2.0),
|
|
tile_grid_size=p.get("clahe_tile_grid_size", (8, 8)),
|
|
)
|
|
|
|
img = threshold(
|
|
img,
|
|
manual_thresh=p.get("threshold_manual"),
|
|
)
|
|
|
|
if p.get("edge_detect", False):
|
|
img = edge_detect(
|
|
img,
|
|
low=p.get("edge_low", 50),
|
|
high=p.get("edge_high", 150),
|
|
)
|
|
|
|
img = morphological_ops(
|
|
img,
|
|
kernel_size=p.get("morph_kernel_size", 3),
|
|
dilate_iterations=p.get("morph_dilate_iterations", 1),
|
|
erode_iterations=p.get("morph_erode_iterations", 1),
|
|
)
|
|
|
|
return img
|