SplatAtlas / tools /build_fig1_separate.py

Upload SplatAtlas benchmark pipeline code

23e73f9 verified 4 months ago

7.75 kB

import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import os
import glob
import json
import numpy as np
from pdf2image import convert_from_path

# =========================================================================
# 配置项
# =========================================================================
SCENE = "bonsai"
VIEW = "00000.png"
OUTPUTS_BASE = "/root/autodl-tmp/SplatAtlas/outputs"
METHODS = ["vanilla_3dgs", "pgsr", "ges", "opti3dgs", "steepgs"]

LEFT_PDF_PATH  = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf"
OUTPUT_DIR = "/root/autodl-tmp/SplatAtlas/tex/figures"

# =========================================================================
# 辅助函数 (构建右图 6宫格 用)
# =========================================================================
def get_valid_dir(method, scene):
    for suffix in ["", "_bak"]:
        d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}"
        if os.path.exists(d) and len(glob.glob(f"{d}/renders_test*")) > 0:
            return d
    return None

def get_metrics(cell_dir):
    try:
        metric_files = sorted(glob.glob(f"{cell_dir}/metrics_test_iter*.json"))
        if not metric_files: return None
        with open(metric_files[-1], 'r') as f:
            data = json.load(f)
            geo = data.get("geometric", {})
            return {
                "opacity": geo.get("alpha_mean", 0.0),
                "aniso": geo.get("gamma_median", 0.0),
                "psnr": data.get("photometric", {}).get("PSNR", 0.0)
            }
    except Exception:
        return None

def add_label(img, title, psnr, opacity=None, aniso=None, is_gt=False):
    """在图片上绘制带有半透明黑底的白字标签"""
    draw = ImageDraw.Draw(img, 'RGBA')
    w, h = img.size
    title_text = title.upper()
    font_size = max(16, int(h * 0.04)) # 字号稍微调大一点点
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size)
    except:
        font = ImageFont.load_default()
    
    padding = 5
    title_bbox = draw.textbbox((0, 0), title_text, font=font)
    tw, th = title_bbox[2] - title_bbox[0], title_bbox[3] - title_bbox[1]
    title_x = (w - tw) // 2
    title_y = 10
    draw.rectangle([title_x - padding, title_y - padding, 
                   title_x + tw + padding, title_y + th + padding], 
                   fill=(0, 0, 0, 180))
    title_color = (255, 255, 255) if is_gt else (100, 255, 100)
    draw.text((title_x, title_y), title_text, font=font, fill=title_color)

    if psnr:
        psnr_text = f"PSNR: {psnr:.2f} dB" if not is_gt else "Reference"
        p_bbox = draw.textbbox((0, 0), psnr_text, font=font)
        pw, ph = p_bbox[2] - p_bbox[0], p_bbox[3] - p_bbox[1]
        px = w - pw - 10
        py = h - ph - 45 # 留出底栏空间
        draw.rectangle([px - padding, py - padding, 
                       px + pw + padding, py + ph + padding], 
                       fill=(0, 0, 0, 180))
        draw.text((px, py), psnr_text, font=font, fill=(255,255,255))
        
    if opacity is not None and aniso is not None:
        bar_h = 35
        draw.rectangle([0, h - bar_h, w, h], fill=(30, 30, 30, 255))
        stat_text = f"opacity μ: {opacity:.3f} | anisotropy μ: {aniso:.3f}" if not is_gt else "opacity μ: 1.000 | anisotropy μ: 1.000"
        s_bbox = draw.textbbox((0, 0), stat_text, font=font)
        sw, sh = s_bbox[2] - s_bbox[0], s_bbox[3] - s_bbox[1]
        draw.text(((w-sw)//2, h - bar_h + (bar_h-sh)//2), stat_text, font=font, fill=(200,200,200))
        
    return img

os.makedirs(OUTPUT_DIR, exist_ok=True)

# =========================================================================
# 处理 Fig 1a: 左图 (Rank Instability) 单独输出
# =========================================================================
print("正在处理 Fig 1a (左图)...")
if not os.path.exists(LEFT_PDF_PATH):
    print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}")
else:
    try:
        pages = convert_from_path(LEFT_PDF_PATH, dpi=300)
        img_left = pages[0]
        
        # 为左图单独创建一个更宽的画板
        fig_a, ax_a = plt.subplots(figsize=(12, 6))
        ax_a.imshow(img_left)
        ax_a.axis('off')
        ax_a.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", 
                      fontsize=20, pad=15, loc='center', fontweight='bold')
        
        out_a_pdf = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.pdf")
        out_a_png = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.png")
        
        fig_a.savefig(out_a_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
        fig_a.savefig(out_a_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
        print(f"✅ Fig 1a 保存成功: {out_a_pdf}")
        plt.close(fig_a)
    except Exception as e:
        print(f"❌ 左图处理失败: {e}")

# =========================================================================
# 处理 Fig 1b: 右图 (6宫格 Render vs Rep) 单独输出
# =========================================================================
print("\n正在处理 Fig 1b (右图 6宫格)...")
images = []
labels = ["GT"] + METHODS

# 1. 提取 GT
gt_dir = get_valid_dir("vanilla_3dgs", SCENE)
if gt_dir:
    gt_path = os.path.join(gt_dir, "gt_test_30000", VIEW)
    if not os.path.exists(gt_path):
        gt_path = os.path.join(gt_dir, "gt_test", VIEW)
    if os.path.exists(gt_path):
        img_gt = Image.open(gt_path).convert("RGBA")
        img_gt = add_label(img_gt, "GT", None, opacity=1.0, aniso=1.0, is_gt=True)
        images.append(img_gt)
    else:
        print(f"❌ 找不到 GT 图像: {gt_path}"); exit(1)
else:
    print("❌ 找不到 vanilla_3dgs 目录以提取 GT"); exit(1)

# 2. 提取 Methods
for method in METHODS:
    cell_dir = get_valid_dir(method, SCENE)
    if not cell_dir:
        images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
        continue
    render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*"))
    if not render_dirs:
         images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
         continue
    img_path = os.path.join(render_dirs[-1], VIEW)
    if not os.path.exists(img_path):
         images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
         continue
    img_m = Image.open(img_path).convert("RGBA")
    metrics = get_metrics(cell_dir)
    if metrics:
        img_m = add_label(img_m, method, metrics['psnr'], metrics['opacity'], metrics['aniso'])
    else:
        img_m = add_label(img_m, method, 0.0)
    images.append(img_m)

# 3. 拼接 6 宫格
iw, ih = images[0].size
padding = 15
grid_w = iw * 3 + padding * 2
grid_h = ih * 2 + padding

img_right = Image.new("RGBA", (grid_w, grid_h), (255,255,255,0)) # 使用透明底色
positions = [
    (0, 0), (iw + padding, 0), ((iw + padding) * 2, 0),
    (0, ih + padding), (iw + padding, ih + padding), ((iw + padding) * 2, ih + padding)
]
for img, pos in zip(images, positions):
    img_right.paste(img, pos)

# 4. 单独制图保存
fig_b, ax_b = plt.subplots(figsize=(14, 8)) # 适合 3x2 的宽幅比例
ax_b.imshow(img_right)
ax_b.axis('off')
ax_b.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation", 
              fontsize=20, pad=15, loc='center', fontweight='bold')

out_b_pdf = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.pdf")
out_b_png = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.png")

fig_b.savefig(out_b_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
fig_b.savefig(out_b_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
print(f"✅ Fig 1b 保存成功: {out_b_pdf}")
plt.close(fig_b)

print("\n=== 图片分离处理完成！===")

	import matplotlib.pyplot as plt
	from PIL import Image, ImageDraw, ImageFont
	import os
	import glob
	import json
	import numpy as np
	from pdf2image import convert_from_path

	# =========================================================================
	# 配置项
	# =========================================================================
	SCENE = "bonsai"
	VIEW = "00000.png"
	OUTPUTS_BASE = "/root/autodl-tmp/SplatAtlas/outputs"
	METHODS = ["vanilla_3dgs", "pgsr", "ges", "opti3dgs", "steepgs"]

	LEFT_PDF_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf"
	OUTPUT_DIR = "/root/autodl-tmp/SplatAtlas/tex/figures"

	# =========================================================================
	# 辅助函数 (构建右图 6宫格用)
	# =========================================================================
	def get_valid_dir(method, scene):
	for suffix in ["", "_bak"]:
	d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}"
	if os.path.exists(d) and len(glob.glob(f"{d}/renders_test*")) > 0:
	return d
	return None

	def get_metrics(cell_dir):
	try:
	metric_files = sorted(glob.glob(f"{cell_dir}/metrics_test_iter*.json"))
	if not metric_files: return None
	with open(metric_files[-1], 'r') as f:
	data = json.load(f)
	geo = data.get("geometric", {})
	return {
	"opacity": geo.get("alpha_mean", 0.0),
	"aniso": geo.get("gamma_median", 0.0),
	"psnr": data.get("photometric", {}).get("PSNR", 0.0)
	}
	except Exception:
	return None

	def add_label(img, title, psnr, opacity=None, aniso=None, is_gt=False):
	"""在图片上绘制带有半透明黑底的白字标签"""
	draw = ImageDraw.Draw(img, 'RGBA')
	w, h = img.size
	title_text = title.upper()
	font_size = max(16, int(h * 0.04)) # 字号稍微调大一点点
	try:
	font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size)
	except:
	font = ImageFont.load_default()

	padding = 5
	title_bbox = draw.textbbox((0, 0), title_text, font=font)
	tw, th = title_bbox[2] - title_bbox[0], title_bbox[3] - title_bbox[1]
	title_x = (w - tw) // 2
	title_y = 10
	draw.rectangle([title_x - padding, title_y - padding,
	title_x + tw + padding, title_y + th + padding],
	fill=(0, 0, 0, 180))
	title_color = (255, 255, 255) if is_gt else (100, 255, 100)
	draw.text((title_x, title_y), title_text, font=font, fill=title_color)

	if psnr:
	psnr_text = f"PSNR: {psnr:.2f} dB" if not is_gt else "Reference"
	p_bbox = draw.textbbox((0, 0), psnr_text, font=font)
	pw, ph = p_bbox[2] - p_bbox[0], p_bbox[3] - p_bbox[1]
	px = w - pw - 10
	py = h - ph - 45 # 留出底栏空间
	draw.rectangle([px - padding, py - padding,
	px + pw + padding, py + ph + padding],
	fill=(0, 0, 0, 180))
	draw.text((px, py), psnr_text, font=font, fill=(255,255,255))

	if opacity is not None and aniso is not None:
	bar_h = 35
	draw.rectangle([0, h - bar_h, w, h], fill=(30, 30, 30, 255))
	stat_text = f"opacity μ: {opacity:.3f} \| anisotropy μ: {aniso:.3f}" if not is_gt else "opacity μ: 1.000 \| anisotropy μ: 1.000"
	s_bbox = draw.textbbox((0, 0), stat_text, font=font)
	sw, sh = s_bbox[2] - s_bbox[0], s_bbox[3] - s_bbox[1]
	draw.text(((w-sw)//2, h - bar_h + (bar_h-sh)//2), stat_text, font=font, fill=(200,200,200))

	return img

	os.makedirs(OUTPUT_DIR, exist_ok=True)

	# =========================================================================
	# 处理 Fig 1a: 左图 (Rank Instability) 单独输出
	# =========================================================================
	print("正在处理 Fig 1a (左图)...")
	if not os.path.exists(LEFT_PDF_PATH):
	print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}")
	else:
	try:
	pages = convert_from_path(LEFT_PDF_PATH, dpi=300)
	img_left = pages[0]

	# 为左图单独创建一个更宽的画板
	fig_a, ax_a = plt.subplots(figsize=(12, 6))
	ax_a.imshow(img_left)
	ax_a.axis('off')
	ax_a.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)",
	fontsize=20, pad=15, loc='center', fontweight='bold')

	out_a_pdf = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.pdf")
	out_a_png = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.png")

	fig_a.savefig(out_a_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
	fig_a.savefig(out_a_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
	print(f"✅ Fig 1a 保存成功: {out_a_pdf}")
	plt.close(fig_a)
	except Exception as e:
	print(f"❌ 左图处理失败: {e}")

	# =========================================================================
	# 处理 Fig 1b: 右图 (6宫格 Render vs Rep) 单独输出
	# =========================================================================
	print("\n正在处理 Fig 1b (右图 6宫格)...")
	images = []
	labels = ["GT"] + METHODS

	# 1. 提取 GT
	gt_dir = get_valid_dir("vanilla_3dgs", SCENE)
	if gt_dir:
	gt_path = os.path.join(gt_dir, "gt_test_30000", VIEW)
	if not os.path.exists(gt_path):
	gt_path = os.path.join(gt_dir, "gt_test", VIEW)
	if os.path.exists(gt_path):
	img_gt = Image.open(gt_path).convert("RGBA")
	img_gt = add_label(img_gt, "GT", None, opacity=1.0, aniso=1.0, is_gt=True)
	images.append(img_gt)
	else:
	print(f"❌ 找不到 GT 图像: {gt_path}"); exit(1)
	else:
	print("❌ 找不到 vanilla_3dgs 目录以提取 GT"); exit(1)

	# 2. 提取 Methods
	for method in METHODS:
	cell_dir = get_valid_dir(method, SCENE)
	if not cell_dir:
	images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
	continue
	render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*"))
	if not render_dirs:
	images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
	continue
	img_path = os.path.join(render_dirs[-1], VIEW)
	if not os.path.exists(img_path):
	images.append(Image.new("RGBA", images[0].size, (50,50,50,255)))
	continue
	img_m = Image.open(img_path).convert("RGBA")
	metrics = get_metrics(cell_dir)
	if metrics:
	img_m = add_label(img_m, method, metrics['psnr'], metrics['opacity'], metrics['aniso'])
	else:
	img_m = add_label(img_m, method, 0.0)
	images.append(img_m)

	# 3. 拼接 6 宫格
	iw, ih = images[0].size
	padding = 15
	grid_w = iw * 3 + padding * 2
	grid_h = ih * 2 + padding

	img_right = Image.new("RGBA", (grid_w, grid_h), (255,255,255,0)) # 使用透明底色
	positions = [
	(0, 0), (iw + padding, 0), ((iw + padding) * 2, 0),
	(0, ih + padding), (iw + padding, ih + padding), ((iw + padding) * 2, ih + padding)
	]
	for img, pos in zip(images, positions):
	img_right.paste(img, pos)

	# 4. 单独制图保存
	fig_b, ax_b = plt.subplots(figsize=(14, 8)) # 适合 3x2 的宽幅比例
	ax_b.imshow(img_right)
	ax_b.axis('off')
	ax_b.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation",
	fontsize=20, pad=15, loc='center', fontweight='bold')

	out_b_pdf = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.pdf")
	out_b_png = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.png")

	fig_b.savefig(out_b_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True)
	fig_b.savefig(out_b_png, dpi=300, bbox_inches='tight', format='png', transparent=False)
	print(f"✅ Fig 1b 保存成功: {out_b_pdf}")
	plt.close(fig_b)

	print("\n=== 图片分离处理完成！===")