| |
| """给某个数据集添加一对「RGB / GT 叠加」样例。 |
| |
| 用法： |
| python tools/add_preview.py <数据集 id> --rgb a_rgb.png --gt a_gt.png \ |
| [--caption "说明"] [--source "GT 来源"] [--name 样例名] [--max-side 960] |
| |
| 图片会缩放到最长边 --max-side 像素，存为 JPEG 放进 media/<id>/， |
| 登记到 data/preview_pool.json，并按现有分类重新抽样。同名样例会被覆盖。 |
| 数据集 id 见 data/datasets.json（页面地址 #ds-<id> 里的那段）。 |
| """ |
| import argparse |
| import json |
| import sys |
| import subprocess |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| def save(src: Path, dst: Path, max_side: int): |
| try: |
| from PIL import Image |
| except ImportError: |
| sys.exit("需要 Pillow：pip install pillow") |
| im = Image.open(src).convert("RGB") |
| w, h = im.size |
| s = max_side / max(w, h) |
| if s < 1: |
| im = im.resize((round(w * s), round(h * s)), Image.LANCZOS) |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| im.save(dst, "JPEG", quality=88, optimize=True) |
| return im.size |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("dataset") |
| ap.add_argument("--rgb", required=True, type=Path) |
| ap.add_argument("--gt", required=True, type=Path) |
| ap.add_argument("--caption", default="") |
| ap.add_argument("--source", default="") |
| ap.add_argument("--name", default="", help="样例名，默认取 RGB 文件名") |
| ap.add_argument("--max-side", type=int, default=960) |
| a = ap.parse_args() |
|
|
| ids = {d["id"] for d in json.loads((ROOT / "data/datasets.json").read_text(encoding="utf-8"))["datasets"]} |
| if a.dataset not in ids: |
| sys.exit(f"未知的数据集 id：{a.dataset}\n可用的 id：{', '.join(sorted(ids))}") |
| for p in (a.rgb, a.gt): |
| if not p.is_file(): |
| sys.exit(f"找不到文件：{p}") |
|
|
| name = a.name or a.rgb.stem.removesuffix("_rgb") |
| rel_rgb = f"media/{a.dataset}/{name}_rgb.jpg" |
| rel_gt = f"media/{a.dataset}/{name}_gt.jpg" |
| width, height = save(a.rgb, ROOT / rel_rgb, a.max_side) |
| save(a.gt, ROOT / rel_gt, a.max_side) |
|
|
| pf = ROOT / "data/preview_pool.json" |
| prev = json.loads(pf.read_text(encoding="utf-8")) if pf.exists() else {} |
| lst = [p for p in prev.get(a.dataset, []) if p.get("rgb") != rel_rgb] |
| lst.append({"id":name,"rgb":rel_rgb,"gt":rel_gt,"modes":[{"label":a.source,"src":rel_gt}], |
| "width":width,"height":height,"caption":a.caption,"source":a.source,"rgb_label":"原始 RGB", |
| "strata":{},"source_page":"","evidence_url":"","license":"","note":""}) |
| prev[a.dataset] = lst |
| pf.write_text(json.dumps(prev, ensure_ascii=False, indent=1) + "\n", encoding="utf-8") |
| subprocess.run([sys.executable, str(ROOT / "tools/import_team_previews.py"), "--per", "12"], check=True) |
| print(f"已登记素材并更新抽样：{a.dataset} / {name}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
