File size: 4,164 Bytes
71dabc7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | #!/usr/bin/env python3
"""把群聊存档与频道导出的图片合并成统一图片池（视觉路线的训练/评测入口）。
两个来源的路径基准不同，这里统一解析成绝对路径并补算 dhash，
跨源再去一次重（同一张广告图被转发到不同 chat 的情况确实存在）。
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import sys
from typing import Iterable
SOURCES = {
"group": ("data/processed/archive/image_clf.jsonl", "data/raw/archive/archive"),
"channel": ("data/processed/channel/images.jsonl", "."),
}
def dhash(path: str, size: int = 8) -> int | None:
from PIL import Image
try:
with Image.open(path) as im:
im = im.convert("L").resize((size + 1, size), Image.Resampling.LANCZOS)
px = list(im.getdata())
except Exception:
return None
bits = 0
for row in range(size):
for col in range(size):
left = px[row * (size + 1) + col]
right = px[row * (size + 1) + col + 1]
bits = (bits << 1) | (1 if left > right else 0)
return bits
def load(manifest: str, base: str, source: str, radius: int) -> Iterable[dict]:
if not os.path.exists(manifest):
print(f"[skip] 缺 {manifest}", file=sys.stderr)
return
with open(manifest, encoding="utf-8") as fh:
for line in fh:
r = json.loads(line)
abs_path = os.path.join(base, r["path"]) if base != "." else os.path.join(os.getcwd(), r["path"])
if not os.path.exists(abs_path):
print(f"[warn] 文件缺失 {abs_path}", file=sys.stderr)
continue
ph = r.get("dhash")
ph = int(ph, 16) if isinstance(ph, str) else dhash(abs_path)
if ph is None:
continue
yield {
"source": source,
"abs_path": os.path.abspath(abs_path),
"chat_id": r.get("chat_id"),
"message_id": r.get("message_id"),
"date": r.get("date"),
"label": r["label"],
"ad_kind": r.get("ad_kind"),
"label_source": r.get("label_source", "judgment"),
"width": r.get("width"),
"height": r.get("height"),
"sha256": r.get("sha256"),
"dhash": f"{ph:016x}",
"_ph": ph,
"text": (r.get("text") or "")[:300],
}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="data/processed/image_pool.jsonl")
ap.add_argument("--phash-radius", type=int, default=4)
args = ap.parse_args()
seen_hash: set[str] = set()
seen_ph: list[tuple[int, dict]] = []
pool, dropped = [], {"exact_dup": 0, "near_dup": 0}
for source, (manifest, base) in SOURCES.items():
n = 0
for r in load(manifest, base, source, args.phash_radius):
n += 1
if r["sha256"] and r["sha256"] in seen_hash:
dropped["exact_dup"] += 1
continue
if any(((r["_ph"] ^ p).bit_count()) <= args.phash_radius for p, _ in seen_ph):
dropped["near_dup"] += 1
continue
if r["sha256"]:
seen_hash.add(r["sha256"])
seen_ph.append((r["_ph"], r))
r.pop("_ph")
pool.append(r)
print(f"[{source}] 读入 {n}")
pool.sort(key=lambda r: (r["label"] != "ad", r["date"] or 0))
with open(args.out, "w", encoding="utf-8") as fh:
for r in pool:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
ads = [r for r in pool if r["label"] == "ad"]
by_source = {}
for r in ads:
by_source[r["source"]] = by_source.get(r["source"], 0) + 1
print(json.dumps({
"pool": len(pool),
"dropped": dropped,
"label_dist": {"ad": len(ads), "clean": len(pool) - len(ads)},
"ad_by_source": by_source,
"out": os.path.abspath(args.out),
}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
sys.exit(main())
|