380 lines
12 KiB
Python
380 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
图片转SVG工具 - 基于 potrace 的高保真描边
|
||
目标：准确还原图片中的黑色部分
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
import tempfile
|
||
|
||
import cv2
|
||
import numpy as np
|
||
|
||
|
||
class ImageToSVG:
|
||
"""图片转SVG转换器 - 使用 potrace 进行高精度矢量化"""
|
||
|
||
def __init__(
|
||
self,
|
||
threshold=None,
|
||
turdsize=0,
|
||
opttolerance=0,
|
||
unit=1,
|
||
optimize_curves=False,
|
||
circle_fit=0.0,
|
||
):
|
||
"""
|
||
初始化转换器
|
||
|
||
Args:
|
||
threshold: 固定阈值（0-255），None 表示使用 Otsu 自动阈值
|
||
turdsize: 抑制噪点的面积阈值（potrace -t）
|
||
opttolerance: 曲线优化容差（potrace -O）
|
||
unit: 坐标量化单位（potrace -u，1 表示不量化）
|
||
optimize_curves: 是否开启曲线优化（文件更小但可能失真）
|
||
circle_fit: 圆拟合误差阈值（相对半径），>0启用圆替代
|
||
"""
|
||
self.threshold = threshold
|
||
self.turdsize = turdsize
|
||
self.opttolerance = opttolerance
|
||
self.unit = unit
|
||
self.optimize_curves = optimize_curves
|
||
self.circle_fit = circle_fit
|
||
|
||
@staticmethod
|
||
def _iter_input_files(input_dir):
|
||
exts = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
|
||
for name in sorted(os.listdir(input_dir)):
|
||
path = os.path.join(input_dir, name)
|
||
if not os.path.isfile(path):
|
||
continue
|
||
if os.path.splitext(name)[1].lower() in exts:
|
||
yield path
|
||
|
||
@staticmethod
|
||
def _ensure_potrace():
|
||
"""确保 potrace 可用"""
|
||
if shutil.which("potrace") is None:
|
||
raise RuntimeError("未找到 potrace，请先安装：brew install potrace")
|
||
|
||
@staticmethod
|
||
def _load_gray(image_path):
|
||
"""加载图片并转换为灰度图（透明通道会按白底合成）"""
|
||
img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
|
||
if img is None:
|
||
raise ValueError(f"无法读取图片: {image_path}")
|
||
|
||
if img.ndim == 3 and img.shape[2] == 4:
|
||
bgr = img[:, :, :3].astype(np.float32)
|
||
alpha = img[:, :, 3:4].astype(np.float32) / 255.0
|
||
bgr = bgr * alpha + 255.0 * (1.0 - alpha)
|
||
gray = cv2.cvtColor(bgr.astype(np.uint8), cv2.COLOR_BGR2GRAY)
|
||
elif img.ndim == 3:
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
else:
|
||
gray = img
|
||
|
||
return gray
|
||
|
||
def _binarize(self, gray):
|
||
"""二值化：黑色为前景"""
|
||
if self.threshold is None:
|
||
_, binary = cv2.threshold(
|
||
gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
|
||
)
|
||
else:
|
||
_, binary = cv2.threshold(gray, self.threshold, 255, cv2.THRESH_BINARY_INV)
|
||
|
||
return (binary > 0).astype(np.uint8)
|
||
|
||
@staticmethod
|
||
def _write_pbm(mask, path):
|
||
"""写入 ASCII PBM（P1）"""
|
||
height, width = mask.shape
|
||
with open(path, "w", encoding="ascii") as f:
|
||
f.write("P1\n")
|
||
f.write(f"{width} {height}\n")
|
||
for y in range(height):
|
||
row = mask[y]
|
||
f.write(" ".join("1" if v else "0" for v in row))
|
||
f.write("\n")
|
||
|
||
def _run_potrace(self, pbm_path, output_path):
|
||
"""调用 potrace 生成 SVG"""
|
||
cmd = [
|
||
"potrace",
|
||
str(pbm_path),
|
||
"-s",
|
||
"-o",
|
||
str(output_path),
|
||
"-u",
|
||
str(self.unit),
|
||
"-t",
|
||
str(self.turdsize),
|
||
"-O",
|
||
str(self.opttolerance),
|
||
]
|
||
if not self.optimize_curves:
|
||
cmd.append("-n")
|
||
subprocess.run(cmd, check=True)
|
||
|
||
@staticmethod
|
||
def _format_number(value):
|
||
text = f"{value:.2f}".rstrip("0").rstrip(".")
|
||
return text if text else "0"
|
||
|
||
def _circle_path(self, cx, cy, radius):
|
||
"""将圆转换为SVG路径（由两段圆弧组成）"""
|
||
cy_s = self._format_number(cy)
|
||
r_s = self._format_number(radius)
|
||
start_x = self._format_number(cx + radius)
|
||
end_x = self._format_number(cx - radius)
|
||
return (
|
||
f"M {start_x} {cy_s} "
|
||
f"A {r_s} {r_s} 0 1 0 {end_x} {cy_s} "
|
||
f"A {r_s} {r_s} 0 1 0 {start_x} {cy_s} Z"
|
||
)
|
||
|
||
@staticmethod
|
||
def _fit_circle(contour):
|
||
"""拟合圆并返回最大相对误差"""
|
||
(cx, cy), radius = cv2.minEnclosingCircle(contour)
|
||
if radius <= 0:
|
||
return None
|
||
points = contour.reshape(-1, 2).astype(np.float32)
|
||
dx = points[:, 0] - cx
|
||
dy = points[:, 1] - cy
|
||
dist = np.sqrt(dx * dx + dy * dy)
|
||
rel_error = np.abs(dist - radius) / radius
|
||
max_rel_error = float(rel_error.max())
|
||
return cx, cy, radius, max_rel_error
|
||
|
||
@staticmethod
|
||
def _collect_component_indices(hierarchy, start_idx):
|
||
indices = []
|
||
stack = [start_idx]
|
||
while stack:
|
||
idx = stack.pop()
|
||
indices.append(idx)
|
||
child = hierarchy[idx][2]
|
||
while child != -1:
|
||
stack.append(child)
|
||
child = hierarchy[child][0]
|
||
return indices
|
||
|
||
def _try_circle_svg(self, mask, output_path):
|
||
if self.circle_fit <= 0:
|
||
return False
|
||
|
||
contours, hierarchy = cv2.findContours(
|
||
mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE
|
||
)
|
||
if not contours or hierarchy is None:
|
||
return False
|
||
|
||
hierarchy = hierarchy[0]
|
||
paths = []
|
||
|
||
for idx, info in enumerate(hierarchy):
|
||
if info[3] != -1:
|
||
continue
|
||
|
||
parent_area = cv2.contourArea(contours[idx])
|
||
if parent_area <= self.turdsize:
|
||
continue
|
||
|
||
component_indices = self._collect_component_indices(hierarchy, idx)
|
||
component_paths = []
|
||
|
||
for contour_idx in component_indices:
|
||
contour = contours[contour_idx]
|
||
result = self._fit_circle(contour)
|
||
if result is None:
|
||
return False
|
||
cx, cy, radius, max_rel_error = result
|
||
if max_rel_error > self.circle_fit:
|
||
return False
|
||
component_paths.append(self._circle_path(cx, cy, radius))
|
||
|
||
if component_paths:
|
||
paths.append(
|
||
f' <path d="{" ".join(component_paths)}" '
|
||
f'fill="#000000" fill-rule="evenodd"/>'
|
||
)
|
||
|
||
if not paths:
|
||
return False
|
||
|
||
height, width = mask.shape
|
||
svg_content = (
|
||
'<?xml version="1.0" encoding="UTF-8" standalone="no"?>\n'
|
||
'<svg xmlns="http://www.w3.org/2000/svg" '
|
||
f'width="{width}" height="{height}" '
|
||
f'viewBox="0 0 {width} {height}">\n'
|
||
f"{chr(10).join(paths)}\n"
|
||
"</svg>\n"
|
||
)
|
||
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
f.write(svg_content)
|
||
|
||
return True
|
||
|
||
def convert_to_svg(self, image_path, output_path=None):
|
||
"""
|
||
将图片转换为SVG（高保真）
|
||
|
||
Args:
|
||
image_path: 输入图片路径
|
||
output_path: 输出SVG路径
|
||
"""
|
||
gray = self._load_gray(image_path)
|
||
mask = self._binarize(gray)
|
||
|
||
if mask.sum() == 0:
|
||
print("未检测到黑色区域，跳过。")
|
||
return None
|
||
|
||
if output_path is None:
|
||
base_name = os.path.splitext(image_path)[0]
|
||
output_path = base_name + ".svg"
|
||
else:
|
||
os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
|
||
|
||
if self._try_circle_svg(mask, output_path):
|
||
file_size = os.path.getsize(output_path)
|
||
print(f"SVG文件已保存: {output_path}")
|
||
print(f"文件大小: {file_size} 字节 ({file_size/1024:.1f} KB)")
|
||
return output_path
|
||
|
||
self._ensure_potrace()
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
pbm_path = os.path.join(tmp_dir, "input.pbm")
|
||
self._write_pbm(mask, pbm_path)
|
||
self._run_potrace(pbm_path, output_path)
|
||
|
||
file_size = os.path.getsize(output_path)
|
||
print(f"SVG文件已保存: {output_path}")
|
||
print(f"文件大小: {file_size} 字节 ({file_size/1024:.1f} KB)")
|
||
return output_path
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(
|
||
description="将图片转换为高保真SVG（使用potrace描边）",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例:
|
||
python pic2svg.py input.png
|
||
python pic2svg.py --indir images
|
||
python pic2svg.py --indir images --outdir output
|
||
python pic2svg.py input.png --outdir output
|
||
python pic2svg.py input.png --turdsize 0 --opttolerance 0 --unit 1
|
||
""",
|
||
)
|
||
|
||
parser.add_argument("input", nargs="?", help="输入图片路径（可选，与 --indir 互斥）")
|
||
parser.add_argument("--indir", help="输入目录（批量转换，与 input 互斥）")
|
||
parser.add_argument(
|
||
"--outdir",
|
||
help="输出目录（自动创建，使用输入文件名.svg）",
|
||
)
|
||
parser.add_argument(
|
||
"--threshold",
|
||
type=int,
|
||
default=None,
|
||
help="固定阈值（0-255），默认使用Otsu自动阈值",
|
||
)
|
||
parser.add_argument(
|
||
"--turdsize",
|
||
type=int,
|
||
default=0,
|
||
help="抑制噪点面积阈值（potrace -t，默认 0）",
|
||
)
|
||
parser.add_argument(
|
||
"--opttolerance",
|
||
type=float,
|
||
default=0.0,
|
||
help="曲线优化容差（potrace -O，默认 0）",
|
||
)
|
||
parser.add_argument(
|
||
"--unit",
|
||
type=int,
|
||
default=1,
|
||
help="坐标量化单位（potrace -u，默认 1）",
|
||
)
|
||
parser.add_argument(
|
||
"--optimize-curves",
|
||
action="store_true",
|
||
help="启用曲线优化（文件更小但可能略失真）",
|
||
)
|
||
parser.add_argument(
|
||
"--circle-fit",
|
||
type=float,
|
||
default=0.0,
|
||
help="圆拟合误差阈值（相对半径），>0启用圆替代，如 0.02",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
if args.input is None and args.indir is None:
|
||
parser.error("必须提供输入文件或 --indir")
|
||
if args.input and args.indir:
|
||
parser.error("input 与 --indir 互斥，请只指定一个")
|
||
if args.indir and not os.path.isdir(args.indir):
|
||
parser.error(f"输入目录不存在: {args.indir}")
|
||
|
||
converter = ImageToSVG(
|
||
threshold=args.threshold,
|
||
turdsize=args.turdsize,
|
||
opttolerance=args.opttolerance,
|
||
unit=args.unit,
|
||
optimize_curves=args.optimize_curves,
|
||
circle_fit=args.circle_fit,
|
||
)
|
||
|
||
if args.indir:
|
||
outdir = args.outdir
|
||
if outdir:
|
||
os.makedirs(outdir, exist_ok=True)
|
||
files = list(converter._iter_input_files(args.indir))
|
||
if not files:
|
||
print("输入目录中未找到可处理的图片文件。")
|
||
return 1
|
||
success = 0
|
||
failed = 0
|
||
for path in files:
|
||
output_path = None
|
||
if outdir:
|
||
base_name = os.path.splitext(os.path.basename(path))[0]
|
||
output_path = os.path.join(outdir, base_name + ".svg")
|
||
try:
|
||
converter.convert_to_svg(path, output_path)
|
||
success += 1
|
||
except Exception as e:
|
||
failed += 1
|
||
print(f"转换失败: {path} -> {e}")
|
||
print(f"批量转换完成: 成功 {success} 个，失败 {failed} 个")
|
||
else:
|
||
output_path = None
|
||
if args.outdir:
|
||
os.makedirs(args.outdir, exist_ok=True)
|
||
base_name = os.path.splitext(os.path.basename(args.input))[0]
|
||
output_path = os.path.join(args.outdir, base_name + ".svg")
|
||
converter.convert_to_svg(args.input, output_path)
|
||
except Exception as e:
|
||
print(f"转换失败: {e}")
|
||
return 1
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|