pqAutomationApp/app/export/image_exporter.py

"""测试结果图表 PNG 导出。"""
import os


def _gamut_refs_for_type(test_type):
    """按测试类型返回需要导出的参考色域列表。"""
    if test_type == "sdr_movie":
        return ["BT.709", "DCI-P3", "BT.2020", "BT.601"]
    return ["BT.709", "DCI-P3", "BT.2020"]


def _gamut_ref_var_attr(test_type):
    if test_type == "screen_module":
        return "screen_gamut_ref_var"
    if test_type == "sdr_movie":
        return "sdr_gamut_ref_var"
    if test_type == "hdr_movie":
        return "hdr_gamut_ref_var"
    return None


# (item_key, fig_attr, filename, allowed_test_types_or_None, default_bbox_inches_auto)
IMAGE_SPECS = [
    ("gamut", "gamut_fig", "色域测试结果.png", None, True),
    ("gamma", "gamma_fig", "Gamma曲线测试结果.png",
     {"screen_module", "sdr_movie"}, True),
    ("eotf", "eotf_fig", "EOTF曲线测试结果.png", {"hdr_movie"}, True),
    ("cct", "cct_fig", "色度一致性测试结果.png", None, True),
    ("contrast", "contrast_fig", "对比度测试结果.png", None, False),
    ("accuracy", "accuracy_fig", "色准测试结果.png",
     {"sdr_movie", "hdr_movie"}, True),
]


def save_result_images(result_dir, current_test_type, selected_items,
                       figure_provider, log, app=None):
    """根据测试类型和已选项将各测试图表保存为 PNG。

    figure_provider: callable(attr_name) -> matplotlib Figure 或 None
    log: callable(msg)
    """
    for item_key, fig_attr, filename, allowed_types, default_bbox in IMAGE_SPECS:
        if item_key not in selected_items:
            continue
        if allowed_types is not None and current_test_type not in allowed_types:
            continue

        # 色域图：按所有可切换参考状态逐个重绘并保存。
        if item_key == "gamut" and app is not None:
            var_attr = _gamut_ref_var_attr(current_test_type)
            ref_var = getattr(app, var_attr, None) if var_attr else None
            rgb_data = None
            coverage = 0.0

            if hasattr(app, "results") and app.results:
                rgb_data = app.results.get_intermediate_data("gamut", "rgb")
                try:
                    coverage = app.results.test_items["gamut"].final_result.get("coverage", 0.0)
                except Exception:
                    coverage = 0.0

            if ref_var is not None and rgb_data and len(rgb_data) >= 3 and hasattr(app, "plot_gamut"):
                original_ref = ref_var.get()
                try:
                    for ref in _gamut_refs_for_type(current_test_type):
                        ref_var.set(ref)
                        app.plot_gamut(rgb_data, coverage, current_test_type)
                        fig = figure_provider(fig_attr)
                        if fig is None:
                            continue
                        per_ref_name = f"色域测试结果_{ref}.png"
                        path = os.path.join(result_dir, per_ref_name)
                        fig.savefig(path, dpi=300)
                        log(f"已保存: {per_ref_name}")
                finally:
                    ref_var.set(original_ref)
                    app.plot_gamut(rgb_data, coverage, current_test_type)
                continue

        fig = figure_provider(fig_attr)
        if fig is None:
            continue
        path = os.path.join(result_dir, filename)
        if default_bbox:
            fig.savefig(path, dpi=300)
        else:
            fig.savefig(path, dpi=300, bbox_inches="tight")
        log(f"已保存: {filename}")