File size: 3,048 Bytes
8aabba6 904b30e 8aabba6 904b30e 8aabba6 904b30e 8aabba6 904b30e 8aabba6 904b30e 8aabba6 904b30e 8aabba6 904b30e 8aabba6 | 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 | #!/usr/bin/env python3
"""给某个数据集添加一对「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()
|