99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
"""
|
|
提示词图像自动切割模块
|
|
将 170x50 的提示图自动切分为 3 个独立的图标小图
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import List, Tuple
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
import io
|
|
|
|
def slice_prompt_image(prompt_img: str | Path | bytes | Image.Image, padding: int = 2) -> List[Tuple[Image.Image, Tuple[int, int, int, int]]]:
|
|
"""
|
|
自动识别并切割提示图中的 3 个字符/图标
|
|
返回: [(PIL_Image, (x1, y1, x2, y2)), ...]
|
|
"""
|
|
if isinstance(prompt_img, Image.Image):
|
|
img = prompt_img.convert("L")
|
|
elif isinstance(prompt_img, bytes):
|
|
img = Image.open(io.BytesIO(prompt_img)).convert("L")
|
|
else:
|
|
img = Image.open(prompt_img).convert("L")
|
|
arr = np.array(img)
|
|
h, w = arr.shape
|
|
|
|
# 背景灰度基准 (通常取左上角像素)
|
|
bg_val = int(arr[0, 0])
|
|
|
|
# 提取暗色笔画像素（排除灰色背景）
|
|
mask = arr < (bg_val - 25)
|
|
|
|
# 水平方向投影，找出连续包含笔画的列区间
|
|
cols_has_dark = mask.any(axis=0)
|
|
|
|
segments: List[Tuple[int, int]] = []
|
|
in_seg = False
|
|
start = 0
|
|
|
|
for i, val in enumerate(cols_has_dark):
|
|
if val and not in_seg:
|
|
in_seg = True
|
|
start = i
|
|
elif not val and in_seg:
|
|
in_seg = False
|
|
if i - start >= 5: # 过滤过窄杂点
|
|
segments.append((start, i))
|
|
if in_seg and (w - start >= 5):
|
|
segments.append((start, w))
|
|
|
|
# 若投影分割不足 3 段（极罕见），按固定 3 等分兜底
|
|
if len(segments) != 3:
|
|
step = w // 3
|
|
segments = [(i * step + 5, (i + 1) * step - 5) for i in range(3)]
|
|
|
|
# 原始彩色图用于切片
|
|
if isinstance(prompt_img, Image.Image):
|
|
orig_img = prompt_img.convert("RGBA")
|
|
elif isinstance(prompt_img, bytes):
|
|
orig_img = Image.open(io.BytesIO(prompt_img)).convert("RGBA")
|
|
else:
|
|
orig_img = Image.open(prompt_img).convert("RGBA")
|
|
crops = []
|
|
|
|
for seg_start, seg_end in segments[:3]:
|
|
# 纵向精确裁剪
|
|
sub_mask = mask[:, seg_start:seg_end]
|
|
rows_has_dark = sub_mask.any(axis=1)
|
|
row_indices = np.where(rows_has_dark)[0]
|
|
|
|
if len(row_indices) > 0:
|
|
top = max(0, int(row_indices[0]) - padding)
|
|
bottom = min(h, int(row_indices[-1]) + 1 + padding)
|
|
else:
|
|
top, bottom = 0, h
|
|
|
|
left = max(0, seg_start - padding)
|
|
right = min(w, seg_end + padding)
|
|
|
|
box = (left, top, right, bottom)
|
|
crop = orig_img.crop(box)
|
|
crops.append((crop, box))
|
|
|
|
return crops
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sample = sys.argv[1] if len(sys.argv) > 1 else "tcaptcha/data/raw/1789135559041_4170_prompt.png"
|
|
if Path(sample).exists():
|
|
results = slice_prompt_image(sample)
|
|
print(f"切片完成，切出 {len(results)} 个目标:")
|
|
for idx, (c, b) in enumerate(results):
|
|
print(f" [{idx+1}] 尺寸: {c.size}, 边界框: {b}")
|
|
else:
|
|
print(f"样本文件不存在: {sample}")
|