pqAutomationApp/algorithm/pq_algorithm.py

import math
import numpy as np
from shapely.geometry import Polygon


# ==================
# 色域覆盖率计算公式
# ==================


def calculate_gamut_coverage(points1, points2):
    """
    计算色域覆盖率（使用 shapely 库）

    Args:
        points1: 测量三角形 [[x1, y1], [x2, y2], [x3, y3]]
        points2: 参考三角形 [[x1, y1], [x2, y2], [x3, y3]]

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）≤ 100%
    """
    try:
        # 创建多边形对象
        poly1 = Polygon(points1)
        poly2 = Polygon(points2)

        # 计算交集
        intersection = poly1.intersection(poly2)

        # 计算面积
        intersection_area = intersection.area
        reference_area = poly2.area

        # 计算覆盖率
        coverage = (
            (intersection_area / reference_area) * 100 if reference_area > 0 else 0.0
        )

        return intersection_area, coverage

    except Exception as e:
        print(f"❌ 计算色域覆盖率失败: {e}")
        return 0.0, 0.0


def calculate_gamut_coverage_DCIP3(points):
    """
    计算 CIE 1931 xy 色度空间下的 DCI-P3 覆盖率

    Args:
        points: [[x1, y1], [x2, y2], [x3, y3]] 测量的 xy 坐标

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）
    """
    # DCI-P3 标准色域坐标（CIE 1931 xy）
    dcip3_xy = [
        [0.680, 0.320],  # Red
        [0.265, 0.690],  # Green
        [0.150, 0.060],  # Blue
    ]

    area, coverage = calculate_gamut_coverage(points, dcip3_xy)
    return area, coverage


def calculate_gamut_coverage_BT709(results):
    """
    计算相对于 BT.709 的色域覆盖率

    Args:
        results: 测量结果列表
                [[x1, y1, lv1, X1, Y1, Z1], ...]
                或 [[x1, y1], [x2, y2], [x3, y3]]

    Returns:
        tuple: (intersection_area, coverage_percentage)
    """
    # BT.709 标准色域三角形（CIE 1931 xy）
    bt709_triangle = [
        [0.6400, 0.3300],  # Red
        [0.3000, 0.6000],  # Green
        [0.1500, 0.0600],  # Blue
    ]

    # 提取测量的 RGB 三个点的 xy 坐标
    if len(results[0]) > 2:
        measured_points = [[r[0], r[1]] for r in results[:3]]
    else:
        measured_points = results[:3]

    area, coverage = calculate_gamut_coverage(measured_points, bt709_triangle)
    return area, coverage


def calculate_gamut_coverage_BT601(results):
    """
    计算相对于 BT.601 的色域覆盖率

    Args:
        results: 测量结果列表

    Returns:
        tuple: (intersection_area, coverage_percentage)
    """
    # BT.601 标准色域三角形（NTSC 版本）
    bt601_triangle = [
        [0.6300, 0.3400],  # Red
        [0.3100, 0.5950],  # Green
        [0.1550, 0.0700],  # Blue
    ]

    # 提取测量点
    if len(results[0]) > 2:
        measured_points = [[r[0], r[1]] for r in results[:3]]
    else:
        measured_points = results[:3]

    area, coverage = calculate_gamut_coverage(measured_points, bt601_triangle)
    return area, coverage


def calculate_gamut_coverage_BT2020(results):
    """
    计算相对于 BT.2020 的色域覆盖率

    Args:
        results: 测量结果列表

    Returns:
        tuple: (intersection_area, coverage_percentage)
    """
    # BT.2020 标准色域三角形
    bt2020_triangle = [
        [0.7080, 0.2920],  # Red
        [0.1700, 0.7970],  # Green
        [0.1310, 0.0460],  # Blue
    ]

    # 提取测量点
    if len(results[0]) > 2:
        measured_points = [[r[0], r[1]] for r in results[:3]]
    else:
        measured_points = results[:3]

    area, coverage = calculate_gamut_coverage(measured_points, bt2020_triangle)
    return area, coverage


def xy_to_uv_1976(x, y):
    """
    CIE 1931 xy → CIE 1976 u'v' 转换（现代标准）

    Args:
        x: CIE 1931 x 坐标
        y: CIE 1931 y 坐标

    Returns:
        (u', v'): CIE 1976 色度坐标
    """
    denom = -2 * x + 12 * y + 3
    if abs(denom) < 1e-10:
        return (0, 0)

    u_prime = (4 * x) / denom
    v_prime = (9 * y) / denom

    return (u_prime, v_prime)


def calculate_uv_gamut_coverage(uv_coords, reference="DCI-P3"):
    """
    计算 CIE 1976 u'v' 色度空间下的色域覆盖率

    Args:
        uv_coords: [[u1, v1], [u2, v2], [u3, v3]] 测量的 u'v' 坐标
        reference: 参考标准 ("DCI-P3", "BT.709", "BT.2020", "BT.601")

    Returns:
        float: 覆盖率百分比
    """
    try:
        # ========== 根据参考标准选择 xy 坐标 ==========
        if reference == "BT.2020":
            ref_xy = [
                (0.7080, 0.2920),  # Red
                (0.1700, 0.7970),  # Green
                (0.1310, 0.0460),  # Blue
            ]
        elif reference == "BT.709":
            ref_xy = [
                (0.6400, 0.3300),  # Red
                (0.3000, 0.6000),  # Green
                (0.1500, 0.0600),  # Blue
            ]
        elif reference == "DCI-P3":
            ref_xy = [
                (0.680, 0.320),  # Red
                (0.265, 0.690),  # Green
                (0.150, 0.060),  # Blue
            ]
        elif reference == "BT.601":
            ref_xy = [
                (0.6300, 0.3400),  # Red
                (0.3100, 0.5950),  # Green
                (0.1550, 0.0700),  # Blue
            ]
        else:
            # 默认使用 DCI-P3
            ref_xy = [
                (0.680, 0.320),
                (0.265, 0.690),
                (0.150, 0.060),
            ]

        # 转换参考标准为 u'v'
        ref_uv = [list(xy_to_uv_1976(x, y)) for x, y in ref_xy]

        # 确保测量坐标格式正确
        measured_uv = [[float(u), float(v)] for u, v in uv_coords]

        # 计算覆盖率
        area, coverage = calculate_gamut_coverage(measured_uv, ref_uv)

        return coverage

    except Exception as e:
        print(f"❌ 计算 u'v' 色域覆盖率失败: {str(e)}")
        import traceback

        traceback.print_exc()
        return 0.0


def calculate_gamut_coverage_DCIP3_uv(uv_points):
    """
    计算 CIE 1976 u'v' 色度空间下的 DCI-P3 覆盖率

    Args:
        uv_points: [[u1, v1], [u2, v2], [u3, v3]] 测量的 u'v' 坐标

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）
    """
    dcip3_xy = [
        (0.680, 0.320),  # Red
        (0.265, 0.690),  # Green
        (0.150, 0.060),  # Blue
    ]

    dcip3_uv = [list(xy_to_uv_1976(x, y)) for x, y in dcip3_xy]
    area, coverage = calculate_gamut_coverage(uv_points, dcip3_uv)
    return area, coverage


def calculate_gamut_coverage_BT709_uv(uv_points):
    """
    计算 CIE 1976 u'v' 色度空间下的 BT.709 覆盖率

    Args:
        uv_points: [[u1, v1], [u2, v2], [u3, v3]] 测量的 u'v' 坐标

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）
    """
    bt709_xy = [
        (0.6400, 0.3300),  # Red
        (0.3000, 0.6000),  # Green
        (0.1500, 0.0600),  # Blue
    ]

    bt709_uv = [list(xy_to_uv_1976(x, y)) for x, y in bt709_xy]
    area, coverage = calculate_gamut_coverage(uv_points, bt709_uv)
    return area, coverage


def calculate_gamut_coverage_BT2020_uv(uv_points):
    """
    计算 CIE 1976 u'v' 色度空间下的 BT.2020 覆盖率

    Args:
        uv_points: [[u1, v1], [u2, v2], [u3, v3]] 测量的 u'v' 坐标

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）
    """
    bt2020_xy = [
        (0.7080, 0.2920),  # Red
        (0.1700, 0.7970),  # Green
        (0.1310, 0.0460),  # Blue
    ]

    bt2020_uv = [list(xy_to_uv_1976(x, y)) for x, y in bt2020_xy]
    area, coverage = calculate_gamut_coverage(uv_points, bt2020_uv)
    return area, coverage


def calculate_gamut_coverage_BT601_uv(uv_points):
    """
    计算 CIE 1976 u'v' 色度空间下的 BT.601 覆盖率

    Args:
        uv_points: [[u1, v1], [u2, v2], [u3, v3]] 测量的 u'v' 坐标

    Returns:
        area: 交集面积
        coverage: 覆盖率（%）
    """
    bt601_xy = [
        (0.6300, 0.3400),  # Red
        (0.3100, 0.5950),  # Green
        (0.1550, 0.0700),  # Blue
    ]

    bt601_uv = [list(xy_to_uv_1976(x, y)) for x, y in bt601_xy]
    area, coverage = calculate_gamut_coverage(uv_points, bt601_uv)
    return area, coverage


# ===============
# Gamma 计算公式
# ===============


def calculate_gamma(results, max_index_fix, pattern_params=None):
    """
    计算 Gamma 值

    Args:
        results: [[x, y, lv, X, Y, Z], ...] 测量结果
        max_index_fix: 最大灰阶索引
        pattern_params: [[R, G, B], ...] 8bit pattern参数，用于计算input_level
                        如果提供，则使用 pattern_value/255 作为底数
                        如果不提供，则使用传统的 灰阶百分比 作为底数

    Returns:
        results_with_gamma: [[x, y, lv, gamma], ...]
    """
    if not results:
        raise ValueError("测量结果为空")

    # 提取亮度值
    luminance_values = [result[2] for result in results]

    if len(luminance_values) < 2:
        raise ValueError(f"数据点数量不足: {len(luminance_values)}")

    # 使用实际的最大亮度值
    max_luminance = max(luminance_values)
    min_luminance = min(luminance_values)

    if max_luminance <= 0:
        raise ValueError(f"最大亮度无效: {max_luminance}")

    # 判断数据排列顺序
    is_descending = luminance_values[0] > luminance_values[-1]

    results_with_gamma = []
    valid_gamma_sum = 0.0
    valid_gamma_count = 0

    for i, result in enumerate(results):
        x, y, lv = result[0], result[1], result[2]

        # 归一化亮度
        L_bar = lv / max_luminance

        # 计算输入灰阶值 input_level
        if pattern_params is not None and i < len(pattern_params):
            # 使用 8bit pattern 数据 / 255 作为底数（22293 Gamma 数据对齐）
            # 取 R 通道值（灰阶图案 R=G=B）
            pattern_value = pattern_params[i][0]
            input_level = pattern_value / 255.0
        else:
            # 传统计算方式：根据数据顺序计算输入灰阶值
            if is_descending:
                # 从亮到暗: i=0 (100%) → i=max (0%)
                input_level = 1.0 - (i / max_index_fix) if max_index_fix > 0 else 0
            else:
                # 从暗到亮: i=0 (0%) → i=max (100%)
                input_level = (i / max_index_fix) if max_index_fix > 0 else 0

        # 计算 Gamma
        gamma = 0.0

        # 跳过极端值（接近 0% 和 100%）
        if input_level <= 0.01 or input_level >= 0.99 or L_bar <= 0.001:
            gamma = 0.0
        else:
            try:
                log_L = math.log(L_bar)
                log_I = math.log(input_level)
                gamma = log_L / log_I

                # 合理性检查
                if 0.5 <= gamma <= 5.0:
                    valid_gamma_sum += gamma
                    valid_gamma_count += 1
                else:
                    gamma = 0.0

            except (ValueError, ZeroDivisionError):
                gamma = 0.0

        results_with_gamma.append(
            [round(x, 7), round(y, 7), round(lv, 7), round(gamma, 4)]
        )

    return results_with_gamma


def calculate_L_bar(results):
    """
    计算归一化亮度列表（用于绘制 Gamma 曲线）

    Args:
        results: 测量结果列表 [[x, y, lv, X, Y, Z], ...]

    Returns:
        L_bar: 归一化亮度列表 [0.0, 0.1, 0.2, ..., 1.0]
    """
    if not results:
        return []

    # 提取亮度值
    luminance_values = [result[2] for result in results]

    # 使用实际的最大亮度值
    max_luminance = max(luminance_values)

    if max_luminance <= 0:
        return [0.0] * len(luminance_values)

    # 归一化
    L_bar = [lv / max_luminance for lv in luminance_values]

    return L_bar


# ============
# 色度坐标提取
# ============


def calculate_cct_from_results(results):
    """
    从测量结果提取色度坐标（xy 坐标）

    ⚠️ 注意：方法名保持不变（历史原因），实际功能是提取色度坐标

    Args:
        results: [[x, y, lv, ...], ...] 测量结果列表

    Returns:
        dict: {
            'x': [x1, x2, x3, ...],
            'y': [y1, y2, y3, ...],
            'lv': [lv1, lv2, lv3, ...]
        }
    """
    x_values = []
    y_values = []
    lv_values = []

    for result in results:
        try:
            if isinstance(result, (list, tuple)) and len(result) >= 2:
                x = float(result[0])
                y = float(result[1])
                x_values.append(x)
                y_values.append(y)

                # 如果有亮度值也保存（可选）
                if len(result) >= 3:
                    lv = float(result[2])
                    lv_values.append(lv)
            else:
                # 数据格式错误时的默认值（D65 白点）
                x_values.append(0.3127)
                y_values.append(0.3290)
                lv_values.append(0.0)

        except Exception as e:
            print(f"⚠️ 提取色度坐标失败: {str(e)}")
            x_values.append(0.3127)
            y_values.append(0.3290)
            lv_values.append(0.0)

    return {"x": x_values, "y": y_values, "lv": lv_values if lv_values else None}


# ========
# 测试代码
# ========
if __name__ == "__main__":
    print("=" * 60)
    print("PQ 算法库测试（使用 shapely 库）")
    print("=" * 60)

    # 测试 1: DCI-P3 覆盖率
    print("\n1. 测试 DCI-P3 覆盖率计算")
    test_points = [[0.680, 0.320], [0.265, 0.690], [0.150, 0.060]]
    area, coverage = calculate_gamut_coverage_DCIP3(test_points)
    print(f"   DCI-P3 覆盖率: {coverage:.2f}% (期望: 100.00%)")

    # 测试 2: BT.709 覆盖率
    print("\n2. 测试 BT.709 覆盖率计算（测量色域超出标准）")
    test_results_bt709 = [
        [0.6637, 0.3236, 41.90],
        [0.3087, 0.6250, 160.29],
        [0.1435, 0.0447, 12.69],
    ]
    area, coverage = calculate_gamut_coverage_BT709(test_results_bt709)
    print(f"   BT.709 覆盖率: {coverage:.2f}% (期望: ~98%)")

    # 测试 3: BT.2020 覆盖率
    print("\n3. 测试 BT.2020 覆盖率计算（测量色域小于标准）")
    test_results_bt2020 = [
        [0.6662, 0.3216, 39.75],
        [0.3096, 0.6282, 155.24],
        [0.1440, 0.0452, 11.96],
    ]
    area, coverage = calculate_gamut_coverage_BT2020(test_results_bt2020)
    print(f"   BT.2020 覆盖率: {coverage:.2f}% (期望: ~60%)")

    # 测试 4: 色度坐标提取
    print("\n4. 测试色度坐标提取")
    test_results = [
        [0.3127, 0.3290, 190.5],
        [0.3125, 0.3288, 150.2],
        [0.3129, 0.3292, 100.8],
    ]
    chromaticity_data = calculate_cct_from_results(test_results)
    print(f"   x 坐标: {chromaticity_data['x']}")
    print(f"   y 坐标: {chromaticity_data['y']}")
    print(f"   亮度值: {chromaticity_data['lv']}")

    # 测试 5: Gamma 计算
    print("\n5. 测试 Gamma 计算")
    test_gamma_results = [
        [0.31, 0.33, 450.0, 0, 0, 0],
        [0.31, 0.33, 350.0, 0, 0, 0],
        [0.31, 0.33, 250.0, 0, 0, 0],
        [0.31, 0.33, 150.0, 0, 0, 0],
        [0.31, 0.33, 50.0, 0, 0, 0],
        [0.31, 0.33, 0.1, 0, 0, 0],
    ]
    gamma_results = calculate_gamma(test_gamma_results, 5)
    for i, result in enumerate(gamma_results):
        print(f"   点{i+1}: Gamma={result[3]:.3f}")

    print("\n" + "=" * 60)
    print("✅ 测试完成！")
    print("=" * 60)