基于边缘的自动白平衡算法_基于边缘的白平衡-优快云博客

本文链接：https://blog.youkuaiyun.com/2401_83885036/article/details/140687973

简介

在图像处理中，自动白平衡（AWB）算法用于调整图像的色彩平衡，使图像中的白色看起来更加真实。常见的AWB方法包括灰度世界假设和完美反射假设。本文介绍一种基于边缘检测的自动白平衡算法，通过在图像边缘两侧选择参考点来进行颜色校正，从而避免大色块的干扰，实现更准确的白平衡调整。

算法步骤

1. 图像边缘检测

首先，我们使用Canny方法对图像进行边缘检测，以找出图像中的边缘。Canny边缘检测是一种常用的边缘检测算法，通过计算图像梯度来检测边缘。

import cv2
import numpy as np
import matplotlib.pyplot as plt

def edge_detection(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    return edges

2. 选择边缘两侧的参考点

在检测到的边缘的两侧各取一定数量的点作为参考点。通过在边缘的两侧各取偏移量为 offset 的点，可以选择更多的参考点参与计算。

def select_edge_points(img, edges, offset=2):
    h, w = edges.shape
    edge_points = []

    for y in range(h):
        for x in range(w):
            if edges[y, x] != 0:
                for dy in [-offset, offset]:
                    for dx in [-offset, offset]:
                        ny, nx = y + dy, x + dx
                        if 0 <= ny < h and 0 <= nx < w:
                            edge_points.append((ny, nx))
    return edge_points

3. 计算参考点的增益

根据选择的参考点，计算RGB通道的平均值，并基于灰度世界假设计算增益。

def calculate_gain_from_edge_points(edge_points, img):
    r_values = []
    g_values = []
    b_values = []

    for y, x in edge_points:
        b, g, r = img[y, x]
        r_values.append(r)
        g_values.append(g)
        b_values.append(b)

    avg_r = np.mean(r_values)
    avg_g = np.mean(g_values)
    avg_b = np.mean(b_values)

    avg_gray = (avg_r + avg_g + avg_b) / 3

    gain_r = avg_gray / avg_r
    gain_g = avg_gray / avg_g
    gain_b = avg_gray / avg_b

    return gain_r, gain_g, gain_b

4. 应用增益进行白平衡调整

应用计算得到的增益对图像进行白平衡调整。

def apply_gain(img, gain_r, gain_g, gain_b):
    img = img.astype(np.float32)
    img[:, :, 2] *= gain_r
    img[:, :, 1] *= gain_g
    img[:, :, 0] *= gain_b
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img

5. 综合实现并标记选择的点

将以上步骤整合在一起，并在原图上标记选择的参考点。

def edge_based_awb(img, offset=2):
    edges = edge_detection(img)
    edge_points = select_edge_points(img, edges, offset)
    gain_r, gain_g, gain_b = calculate_gain_from_edge_points(edge_points, img)
    awb_img = apply_gain(img, gain_r, gain_g, gain_b)
    
    # 标记选择的点
    img_with_points = img.copy()
    for y, x in edge_points:
        cv2.circle(img_with_points, (x, y), 1, (0, 255, 0), -1)
    
    return awb_img, edges, edge_points, img_with_points

# 读取图像
file_path = 'path/to/your/image.jpg'
img = cv2.imread(file_path)

# 应用基于边缘的AWB
awb_img, edges, edge_points, img_with_points = edge_based_awb(img)

# 显示原始图像和校正后的图像
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
axes[0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
axes[0].set_title('Original Image')
axes[1].imshow(cv2.cvtColor(awb_img, cv2.COLOR_BGR2RGB))
axes[1].set_title('Edge-based AWB Image')
axes[2].imshow(cv2.cvtColor(img_with_points, cv2.COLOR_BGR2RGB))
axes[2].set_title('Selected Edge Points')
plt.show()

# 显示边缘图像
plt.figure()
plt.imshow(edges, cmap='gray')
plt.title('Edge Detection')
plt.show()