85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
"""
|
||
Assemble panorama strips into a single comparison grid image.
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import glob
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
from config import MODEL_ORDER
|
||
|
||
|
||
LABEL_W = 100
|
||
STEP_LABEL_H = 40
|
||
PAD = 2
|
||
|
||
|
||
def find_panorama(base_dir, tag, seed):
|
||
prefix_map = {"Q2_K": "q2k", "Q3_K_S": "q3ks", "Q4_0": "q4_0",
|
||
"Q4_1": "q4_1", "Q5_0": "q5_0", "Q5_1": "q5_1",
|
||
"Q6_K": "q6k", "Q8_0": "q8_0"}
|
||
prefix = prefix_map.get(tag, tag.lower())
|
||
pattern = os.path.join(base_dir, f"panorama_{prefix}_s{seed}_*.png")
|
||
files = sorted(glob.glob(pattern))
|
||
return files[-1] if files else None
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--panorama-dir", required=True)
|
||
parser.add_argument("--output-dir", required=True)
|
||
parser.add_argument("--seed", type=int, default=42)
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
||
rows = []
|
||
for tag in MODEL_ORDER:
|
||
path = find_panorama(args.panorama_dir, tag, args.seed)
|
||
if path:
|
||
rows.append((tag, Image.open(path)))
|
||
print(f" [OK] {tag}")
|
||
else:
|
||
print(f" [--] {tag}: not found")
|
||
|
||
if not rows:
|
||
print("No panoramas found.")
|
||
return
|
||
|
||
pw, ph = rows[0][1].size
|
||
n_steps = 20
|
||
step_w = pw // n_steps
|
||
|
||
grid_w = LABEL_W + pw + PAD
|
||
grid_h = STEP_LABEL_H + (ph + PAD) * len(rows)
|
||
grid = Image.new("RGB", (grid_w, grid_h), (20, 20, 20))
|
||
draw = ImageDraw.Draw(grid)
|
||
|
||
try:
|
||
font = ImageFont.truetype("arial.ttf", 20)
|
||
font_sm = ImageFont.truetype("arial.ttf", 14)
|
||
except:
|
||
font = ImageFont.load_default()
|
||
font_sm = font
|
||
|
||
for j in range(n_steps):
|
||
x = LABEL_W + j * step_w + step_w // 2 - 8
|
||
draw.text((x, 10), str(j+1), fill="white", font=font_sm)
|
||
|
||
for i, (tag, pano) in enumerate(rows):
|
||
y = STEP_LABEL_H + (ph + PAD) * i
|
||
draw.text((8, y + ph // 2 - 10), tag, fill=(255, 220, 80), font=font)
|
||
grid.paste(pano, (LABEL_W, y))
|
||
|
||
out_path = os.path.join(args.output_dir, "comparison_grid.png")
|
||
grid.save(out_path, quality=95)
|
||
print(f"\nSaved: {out_path} ({grid_w}×{grid_h})")
|
||
|
||
# Half-size version
|
||
small = grid.resize((grid_w // 2, grid_h // 2), Image.LANCZOS)
|
||
small.save(out_path.replace(".png", "_half.png"), quality=95)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|