300 lines
9.4 KiB
Python
300 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}
|
|
|
|
|
|
def read_image(path: Path) -> np.ndarray:
|
|
raw = np.fromfile(str(path), dtype=np.uint8)
|
|
image = cv2.imdecode(raw, cv2.IMREAD_COLOR)
|
|
if image is None:
|
|
raise ValueError(f"Cannot read image: {path}")
|
|
return image
|
|
|
|
|
|
def write_image(path: Path, image: np.ndarray) -> None:
|
|
ok, encoded = cv2.imencode(path.suffix, image)
|
|
if not ok:
|
|
raise ValueError(f"Cannot encode image: {path}")
|
|
encoded.tofile(str(path))
|
|
|
|
|
|
def binarize_foreground(gray: np.ndarray) -> np.ndarray:
|
|
threshold, _ = cv2.threshold(
|
|
gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
|
|
)
|
|
gray_min = int(gray.min())
|
|
gray_max = int(gray.max())
|
|
if threshold <= gray_min or threshold >= gray_max:
|
|
threshold = (gray_min + gray_max) / 2.0
|
|
|
|
border = np.concatenate([gray[0, :], gray[-1, :], gray[:, 0], gray[:, -1]])
|
|
background_is_light = np.median(border) >= 128
|
|
if background_is_light:
|
|
return (gray < threshold).astype(np.uint8) * 255
|
|
return (gray > threshold).astype(np.uint8) * 255
|
|
|
|
|
|
def foreground_bbox(mask: np.ndarray) -> tuple[int, int, int, int]:
|
|
ys, xs = np.where(mask > 0)
|
|
if xs.size == 0:
|
|
height, width = mask.shape
|
|
return 0, 0, width, height
|
|
return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)
|
|
|
|
|
|
def remove_tiny_components(mask: np.ndarray, min_area: int) -> np.ndarray:
|
|
if min_area <= 0:
|
|
return mask
|
|
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
|
|
mask, connectivity=8
|
|
)
|
|
cleaned = np.zeros_like(mask)
|
|
for label in range(1, num_labels):
|
|
area = int(stats[label, cv2.CC_STAT_AREA])
|
|
if area >= min_area:
|
|
cleaned[labels == label] = 255
|
|
return cleaned
|
|
|
|
|
|
def remove_isolated_specks(
|
|
mask: np.ndarray,
|
|
max_area: int,
|
|
support_radius: int,
|
|
) -> np.ndarray:
|
|
if max_area <= 0 or support_radius <= 0:
|
|
return mask
|
|
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
|
|
mask, connectivity=8
|
|
)
|
|
support = np.zeros_like(mask)
|
|
candidate_labels: list[int] = []
|
|
|
|
for label in range(1, num_labels):
|
|
area = int(stats[label, cv2.CC_STAT_AREA])
|
|
if area <= max_area:
|
|
candidate_labels.append(label)
|
|
else:
|
|
support[labels == label] = 255
|
|
|
|
if not candidate_labels or not np.any(support):
|
|
return mask
|
|
|
|
kernel_size = support_radius * 2 + 1
|
|
kernel = cv2.getStructuringElement(
|
|
cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)
|
|
)
|
|
nearby_support = cv2.dilate(support, kernel)
|
|
cleaned = mask.copy()
|
|
|
|
for label in candidate_labels:
|
|
component = labels == label
|
|
if not np.any(nearby_support[component]):
|
|
cleaned[component] = 0
|
|
|
|
return cleaned
|
|
|
|
|
|
def remove_border_hairlines(mask: np.ndarray, max_thickness: int = 8) -> np.ndarray:
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
|
|
mask, connectivity=8
|
|
)
|
|
height, width = mask.shape
|
|
cleaned = mask.copy()
|
|
border_margin = max(4, int(round(min(width, height) * 0.02)))
|
|
|
|
for label in range(1, num_labels):
|
|
x = int(stats[label, cv2.CC_STAT_LEFT])
|
|
y = int(stats[label, cv2.CC_STAT_TOP])
|
|
comp_width = int(stats[label, cv2.CC_STAT_WIDTH])
|
|
comp_height = int(stats[label, cv2.CC_STAT_HEIGHT])
|
|
near_border = (
|
|
x <= border_margin
|
|
or y <= border_margin
|
|
or x + comp_width >= width - border_margin
|
|
or y + comp_height >= height - border_margin
|
|
)
|
|
vertical_hairline = comp_width <= max_thickness and comp_height >= 24
|
|
horizontal_hairline = comp_height <= max_thickness and comp_width >= 24
|
|
if near_border and (vertical_hairline or horizontal_hairline):
|
|
cleaned[labels == label] = 0
|
|
|
|
return cleaned
|
|
|
|
|
|
def fill_small_holes(mask: np.ndarray, max_area: int) -> np.ndarray:
|
|
if max_area <= 0:
|
|
return mask
|
|
|
|
foreground = mask > 0
|
|
inverse = (~foreground).astype(np.uint8)
|
|
flood_mask = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)
|
|
cv2.floodFill(inverse, flood_mask, (0, 0), 2)
|
|
|
|
holes = (inverse == 1).astype(np.uint8)
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
|
|
holes, connectivity=8
|
|
)
|
|
|
|
filled = foreground.copy()
|
|
for label in range(1, num_labels):
|
|
area = int(stats[label, cv2.CC_STAT_AREA])
|
|
width = int(stats[label, cv2.CC_STAT_WIDTH])
|
|
height = int(stats[label, cv2.CC_STAT_HEIGHT])
|
|
if area <= max_area and max(width, height) <= max(3, int(max_area**0.5) * 4):
|
|
filled[labels == label] = True
|
|
|
|
return filled.astype(np.uint8) * 255
|
|
|
|
|
|
def smooth_edges_once(mask: np.ndarray) -> np.ndarray:
|
|
foreground = mask > 0
|
|
kernel = np.ones((3, 3), dtype=np.uint8)
|
|
neighbors = cv2.filter2D(
|
|
foreground.astype(np.uint8),
|
|
ddepth=-1,
|
|
kernel=kernel,
|
|
borderType=cv2.BORDER_CONSTANT,
|
|
)
|
|
smoothed = (foreground & (neighbors >= 3)) | (~foreground & (neighbors >= 7))
|
|
return smoothed.astype(np.uint8) * 255
|
|
|
|
|
|
def clean_mask(
|
|
mask: np.ndarray,
|
|
min_component_area: int,
|
|
max_hole_area: int,
|
|
max_isolated_speck_area: int,
|
|
smooth: bool = True,
|
|
) -> np.ndarray:
|
|
x1, y1, x2, y2 = foreground_bbox(mask)
|
|
support_radius = max(12, min(42, int(round(min(x2 - x1, y2 - y1) * 0.035))))
|
|
cleaned = remove_tiny_components(mask, min_component_area)
|
|
cleaned = remove_isolated_specks(
|
|
cleaned,
|
|
max_area=max_isolated_speck_area,
|
|
support_radius=support_radius,
|
|
)
|
|
cleaned = remove_border_hairlines(cleaned)
|
|
cleaned = fill_small_holes(cleaned, max_hole_area)
|
|
if smooth:
|
|
cleaned = smooth_edges_once(cleaned)
|
|
cleaned = remove_isolated_specks(
|
|
cleaned,
|
|
max_area=max_isolated_speck_area,
|
|
support_radius=support_radius,
|
|
)
|
|
cleaned = remove_border_hairlines(cleaned)
|
|
cleaned = fill_small_holes(cleaned, max_hole_area)
|
|
return remove_tiny_components(cleaned, min_component_area)
|
|
|
|
|
|
def adaptive_clean_params(mask: np.ndarray) -> tuple[int, int, int]:
|
|
x1, y1, x2, y2 = foreground_bbox(mask)
|
|
bbox_area = max(1, (x2 - x1) * (y2 - y1))
|
|
min_component_area = max(4, min(24, int(round(bbox_area * 0.000018))))
|
|
max_hole_area = max(12, min(90, int(round(bbox_area * 0.00008))))
|
|
max_isolated_speck_area = max(
|
|
min_component_area + 4,
|
|
min(48, int(round(bbox_area * 0.000055))),
|
|
)
|
|
return min_component_area, max_hole_area, max_isolated_speck_area
|
|
|
|
|
|
def mask_to_display(mask: np.ndarray) -> np.ndarray:
|
|
return np.where(mask > 0, 0, 255).astype(np.uint8)
|
|
|
|
|
|
def process_image(
|
|
image_path: Path,
|
|
output_dir: Path,
|
|
min_component_area: int | None,
|
|
max_hole_area: int | None,
|
|
max_isolated_speck_area: int | None,
|
|
smooth: bool,
|
|
) -> None:
|
|
original = read_image(image_path)
|
|
gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
|
|
binary = binarize_foreground(gray)
|
|
(
|
|
adaptive_min_area,
|
|
adaptive_max_hole_area,
|
|
adaptive_max_isolated_speck_area,
|
|
) = adaptive_clean_params(binary)
|
|
cleaned = clean_mask(
|
|
binary,
|
|
min_component_area=(
|
|
adaptive_min_area if min_component_area is None else min_component_area
|
|
),
|
|
max_hole_area=(
|
|
adaptive_max_hole_area if max_hole_area is None else max_hole_area
|
|
),
|
|
max_isolated_speck_area=(
|
|
adaptive_max_isolated_speck_area
|
|
if max_isolated_speck_area is None
|
|
else max_isolated_speck_area
|
|
),
|
|
smooth=smooth,
|
|
)
|
|
write_image(output_dir / f"{image_path.stem}_cleaned.png", mask_to_display(cleaned))
|
|
|
|
|
|
def iter_images(input_dir: Path) -> list[Path]:
|
|
return sorted(
|
|
path
|
|
for path in input_dir.iterdir()
|
|
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
|
|
)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Gently clean complex seal-script characters."
|
|
)
|
|
parser.add_argument(
|
|
"--input-dir",
|
|
type=Path,
|
|
default=Path("PreProcessing_Common/original_characters"),
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=Path,
|
|
default=Path("PreProcessing_Common/cleaned_characters"),
|
|
)
|
|
parser.add_argument("--min-component-area", type=int, default=None)
|
|
parser.add_argument("--max-hole-area", type=int, default=None)
|
|
parser.add_argument("--max-isolated-speck-area", type=int, default=None)
|
|
parser.add_argument("--no-smooth", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
image_paths = iter_images(args.input_dir)
|
|
if not image_paths:
|
|
raise SystemExit(f"No images found in {args.input_dir}")
|
|
|
|
for image_path in image_paths:
|
|
process_image(
|
|
image_path=image_path,
|
|
output_dir=args.output_dir,
|
|
min_component_area=args.min_component_area,
|
|
max_hole_area=args.max_hole_area,
|
|
max_isolated_speck_area=args.max_isolated_speck_area,
|
|
smooth=not args.no_smooth,
|
|
)
|
|
print(f"processed: {image_path.name}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|