satlink: add RS-Tokenizer backend (256x18b, ~114x, ~41dB) alongside One-D-Piece; pixel path only

61a093f verified 2 months ago

2.87 kB

"""The shortest end-to-end call of satlink, as an application would write it.
    python -m satlink.examples.quickstart --image satlink/assets/grid3x3_600x600.png --out /tmp/qs
Runs both halves in one process for readability. In deployment they are two processes on two
machines and the only thing that crosses is `codes.json` (or the packed bytes) -- which is the
point: nothing here reaches across the split except that file.
"""
import argparse
import os
import sys

import satlink


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True, help="PNG or BMP")
    ap.add_argument("--out", default="/tmp/satlink_quickstart")
    ap.add_argument("--sat-weights", default="satlink/weights/sat_encoder.pt")
    ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt")
    ap.add_argument("--model", default=None, help="RSThinker dir; omit to skip the LLM step")
    ap.add_argument("--question", default="describe the image")
    ap.add_argument("--batch", type=int, default=64, help="crops per encoder call")
    ap.add_argument("--dtype", default="float16")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)

    # ---------------------------------------------------------------- 星上
    # The profile is a deployment constant; here it is read off the file so the example runs on
    # any input. In flight the size is known ahead of time and this lookup does not happen.
    from satlink.imageio import StripReader
    with StripReader(a.image) as r:
        w, h = r.width, r.height
    sat = satlink.Satellite(a.sat_weights, w, h, dtype=a.dtype)
    frame = sat.encode_file(a.image, batch=a.batch)
    codes = os.path.join(a.out, "codes.json")
    sat.save(frame, codes)
    wire = sat.to_wire(frame)
    print(f"星上  {w}x{h} -> {sat.n_crops} 个切片 x {frame.profile.n_tok} token "
          f"= {len(wire)} 字节下行，压缩比 {sat.compression_ratio:.1f}x")
    print(f"      落盘 {codes}")

    # ---------------------------------------------------------------- 链路
    # 下行只有码字。帧尺寸和 token 几何都不在流里，是两端约定的常量，所以这里要显式给出。
    # 这里直接把 profile 传过去，它把四个常量一起带上——只对齐 n_tok 而漏掉 bits，
    # 正是会解出错误长度的那种失误。
    frame = satlink.frame_from_wire(wire, profile=frame.profile)

    # ---------------------------------------------------------------- 地面
    recon = os.path.join(a.out, "recon.png")
    satlink.Ground(a.decoder).reconstruct(frame).save(recon)
    print(f"地面  重建 -> {recon}")

    if a.model:
        vlm = satlink.GroundVLM(a.model, decoder=a.decoder)
        print(f"      Q: {a.question}\n      A: {vlm.ask(frame, a.question, max_new=128)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())

	"""The shortest end-to-end call of satlink, as an application would write it.

	python -m satlink.examples.quickstart --image satlink/assets/grid3x3_600x600.png --out /tmp/qs

	Runs both halves in one process for readability. In deployment they are two processes on two
	machines and the only thing that crosses is `codes.json` (or the packed bytes) -- which is the
	point: nothing here reaches across the split except that file.
	"""
	import argparse
	import os
	import sys

	import satlink


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--image", required=True, help="PNG or BMP")
	ap.add_argument("--out", default="/tmp/satlink_quickstart")
	ap.add_argument("--sat-weights", default="satlink/weights/sat_encoder.pt")
	ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt")
	ap.add_argument("--model", default=None, help="RSThinker dir; omit to skip the LLM step")
	ap.add_argument("--question", default="describe the image")
	ap.add_argument("--batch", type=int, default=64, help="crops per encoder call")
	ap.add_argument("--dtype", default="float16")
	a = ap.parse_args()
	os.makedirs(a.out, exist_ok=True)

	# ---------------------------------------------------------------- 星上
	# The profile is a deployment constant; here it is read off the file so the example runs on
	# any input. In flight the size is known ahead of time and this lookup does not happen.
	from satlink.imageio import StripReader
	with StripReader(a.image) as r:
	w, h = r.width, r.height
	sat = satlink.Satellite(a.sat_weights, w, h, dtype=a.dtype)
	frame = sat.encode_file(a.image, batch=a.batch)
	codes = os.path.join(a.out, "codes.json")
	sat.save(frame, codes)
	wire = sat.to_wire(frame)
	print(f"星上 {w}x{h} -> {sat.n_crops} 个切片 x {frame.profile.n_tok} token "
	f"= {len(wire)} 字节下行，压缩比 {sat.compression_ratio:.1f}x")
	print(f" 落盘 {codes}")

	# ---------------------------------------------------------------- 链路
	# 下行只有码字。帧尺寸和 token 几何都不在流里，是两端约定的常量，所以这里要显式给出。
	# 这里直接把 profile 传过去，它把四个常量一起带上——只对齐 n_tok 而漏掉 bits，
	# 正是会解出错误长度的那种失误。
	frame = satlink.frame_from_wire(wire, profile=frame.profile)

	# ---------------------------------------------------------------- 地面
	recon = os.path.join(a.out, "recon.png")
	satlink.Ground(a.decoder).reconstruct(frame).save(recon)
	print(f"地面重建 -> {recon}")

	if a.model:
	vlm = satlink.GroundVLM(a.model, decoder=a.decoder)
	print(f" Q: {a.question}\n A: {vlm.ask(frame, a.question, max_new=128)}")
	return 0


	if __name__ == "__main__":
	sys.exit(main())