103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
"""关键帧字幕区二值掩码哈希聚类 + 拼字幕长图，供多模态批量认字。
|
||
|
|
用法：python cluster_subs.py --frames <frames目录> --outdir <长图输出目录> [--thresh 8] [--crop 0.15] [--rows 10]
|
||
|
|
生成：<outdir>/sheet_001.jpg ... 每张含 rows 句字幕区（原分辨率竖排+时间戳），及 reps.txt 映射。
|
||
|
|
"""
|
||
|
|
import argparse, glob, os, re
|
||
|
|
from PIL import Image, ImageDraw, ImageFont
|
||
|
|
|
||
|
|
|
||
|
|
def crop_subtitle(im, crop):
|
||
|
|
w, h = im.size
|
||
|
|
sh = max(1, int(h * crop))
|
||
|
|
return im.crop((0, h - sh, w, h))
|
||
|
|
|
||
|
|
|
||
|
|
def bin_hash(img, size=24):
|
||
|
|
"""固定高阈值只抓亮色字幕文字（背景归零）→ 同句字幕跨帧掩码稳定可合并。"""
|
||
|
|
g = img.convert("L").resize((size + 1, size), Image.LANCZOS)
|
||
|
|
px = list(g.getdata())
|
||
|
|
thr = 160 # 字幕文字通常近白(>200)，背景暗(<100)，固定阈值隔离文字
|
||
|
|
bits = []
|
||
|
|
for y in range(size):
|
||
|
|
row = px[y * (size + 1):(y + 1) * (size + 1)]
|
||
|
|
for x in range(size):
|
||
|
|
bits.append(1 if row[x] > thr else 0)
|
||
|
|
return bits
|
||
|
|
|
||
|
|
|
||
|
|
def hamming(a, b):
|
||
|
|
return sum(1 for i in range(len(a)) if a[i] != b[i])
|
||
|
|
|
||
|
|
|
||
|
|
def ts_from_name(name):
|
||
|
|
m = re.search(r"_([0-9]+\.[0-9]+)s\.jpg$", name)
|
||
|
|
return float(m.group(1)) if m else 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--frames", required=True)
|
||
|
|
ap.add_argument("--outdir", required=True)
|
||
|
|
ap.add_argument("--thresh", type=int, default=25)
|
||
|
|
ap.add_argument("--crop", type=float, default=0.15)
|
||
|
|
ap.add_argument("--rows", type=int, default=10, help="每张长图放几句")
|
||
|
|
args = ap.parse_args()
|
||
|
|
|
||
|
|
os.makedirs(args.outdir, exist_ok=True)
|
||
|
|
files = sorted(glob.glob(os.path.join(args.frames, "*.jpg")))
|
||
|
|
print("总帧数: %d" % len(files))
|
||
|
|
|
||
|
|
# 聚类
|
||
|
|
clusters = []
|
||
|
|
for f in files:
|
||
|
|
sub = crop_subtitle(Image.open(f), args.crop)
|
||
|
|
hsh = bin_hash(sub)
|
||
|
|
best, bestd = None, 999
|
||
|
|
for ci, c in enumerate(clusters):
|
||
|
|
d = hamming(hsh, c["hash"])
|
||
|
|
if d < bestd:
|
||
|
|
bestd, best = d, ci
|
||
|
|
ts = ts_from_name(os.path.basename(f))
|
||
|
|
if best is not None and bestd <= args.thresh:
|
||
|
|
clusters[best]["items"].append((f, ts))
|
||
|
|
else:
|
||
|
|
clusters.append({"hash": hsh, "items": [(f, ts)]})
|
||
|
|
|
||
|
|
reps = []
|
||
|
|
for c in clusters:
|
||
|
|
it = sorted(c["items"], key=lambda x: x[1])
|
||
|
|
mid = it[len(it) // 2]
|
||
|
|
reps.append((mid[0], mid[1], len(it)))
|
||
|
|
reps.sort(key=lambda x: x[1])
|
||
|
|
print("聚类数(代表帧): %d" % len(reps))
|
||
|
|
|
||
|
|
# 拼长图
|
||
|
|
try:
|
||
|
|
font = ImageFont.load_default()
|
||
|
|
except Exception:
|
||
|
|
font = None
|
||
|
|
sheets = [reps[i:i + args.rows] for i in range(0, len(reps), args.rows)]
|
||
|
|
map_lines = []
|
||
|
|
for si, grp in enumerate(sheets, 1):
|
||
|
|
crops = [crop_subtitle(Image.open(p), args.crop) for p, _, _ in grp]
|
||
|
|
wmax = max(c.width for c in crops)
|
||
|
|
ch = max(c.height for c in crops)
|
||
|
|
label_h = 28
|
||
|
|
canvas = Image.new("RGB", (wmax, (ch + label_h) * len(grp)), (255, 255, 255))
|
||
|
|
d = ImageDraw.Draw(canvas)
|
||
|
|
yy = 0
|
||
|
|
for (p, ts, n), c in zip(grp, crops):
|
||
|
|
d.rectangle([0, yy, wmax, yy + label_h], fill=(230, 230, 230))
|
||
|
|
d.text((6, yy + 6), "%.3fs (同句帧数=%d)" % (ts, n), fill=(0, 0, 0), font=font)
|
||
|
|
canvas.paste(c, (0, yy + label_h))
|
||
|
|
yy += ch + label_h
|
||
|
|
map_lines.append("%03d\t%.3f" % (si, ts))
|
||
|
|
canvas.save(os.path.join(args.outdir, "sheet_%03d.jpg" % si))
|
||
|
|
with open(os.path.join(args.outdir, "reps.txt"), "w", encoding="utf-8") as fo:
|
||
|
|
fo.write("\n".join(map_lines))
|
||
|
|
print("长图张数: %d → %s" % (len(sheets), args.outdir))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|