92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
|
import numpy as np
|
|||
|
|
from cv2 import xphoto
|
|||
|
|
import cv2
|
|||
|
|
from PIL import Image, ImageOps
|
|||
|
|
from matplotlib.patches import Rectangle
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
|
|||
|
|
|
|||
|
|
def learnWb(img):
|
|||
|
|
wb = xphoto.createLearningBasedWB('color_balance_model.yml')
|
|||
|
|
# 读取图像
|
|||
|
|
# image = cv2.imread('IMG_3575.png')
|
|||
|
|
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|||
|
|
# 应用白平衡
|
|||
|
|
white_balanced_image = wb.balanceWhite(img)
|
|||
|
|
# 显示或保存结果
|
|||
|
|
cv2.imshow('White Balanced Image', white_balanced_image)
|
|||
|
|
cv2.waitKey(0)
|
|||
|
|
cv2.destroyAllWindows()
|
|||
|
|
|
|||
|
|
def ground_truth(image, img_patch, mode='mean'):
|
|||
|
|
# 使用均值或最大值来进行颜色校正
|
|||
|
|
if mode == 'mean':
|
|||
|
|
image_gt = ((image * (img_patch.mean() / image.mean(axis=(0, 1)))).clip(0, 255).astype(int))
|
|||
|
|
elif mode == 'max':
|
|||
|
|
image_gt = ((image * 1.0 / img_patch.max(axis=(0, 1))).clip(0, 1))
|
|||
|
|
else:
|
|||
|
|
raise ValueError("Invalid mode. Use 'mean' or 'max'.")
|
|||
|
|
|
|||
|
|
# 绘制原始图像和Ground Truth校正后的对比图
|
|||
|
|
fig, ax = plt.subplots(1, 2, figsize=(14, 10))
|
|||
|
|
ax[0].imshow(image)
|
|||
|
|
ax[0].set_title('原始图像')
|
|||
|
|
ax[0].axis('off')
|
|||
|
|
|
|||
|
|
ax[1].imshow(image_gt)
|
|||
|
|
ax[1].set_title('Ground Truth校正后的图像')
|
|||
|
|
ax[1].axis('off')
|
|||
|
|
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gray_world(image_path):
|
|||
|
|
# 读取图像
|
|||
|
|
image = cv2.imread(image_path)
|
|||
|
|
|
|||
|
|
# 转换图像颜色通道顺序（OpenCV使用BGR，而matplotlib使用RGB）
|
|||
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|||
|
|
|
|||
|
|
# 计算图像的平均亮度
|
|||
|
|
avg_brightness = image.mean()
|
|||
|
|
|
|||
|
|
# 计算每个颜色通道的平均亮度
|
|||
|
|
avg_channel_brightness = image.mean(axis=(0, 1))
|
|||
|
|
|
|||
|
|
# 计算亮度比例，并相应地调整每个颜色通道
|
|||
|
|
image_grayworld = ((image * (avg_brightness / avg_channel_brightness)).clip(0, 255).astype(int))
|
|||
|
|
|
|||
|
|
# 绘制原始图像和灰色世界校正后的对比图
|
|||
|
|
fig, ax = plt.subplots(1, 2, figsize=(14, 10))
|
|||
|
|
ax[0].imshow(image)
|
|||
|
|
ax[0].set_title('原始图像')
|
|||
|
|
ax[0].axis('off')
|
|||
|
|
|
|||
|
|
ax[1].imshow(image_grayworld)
|
|||
|
|
ax[1].set_title('灰色世界校正后的图像')
|
|||
|
|
ax[1].axis('off')
|
|||
|
|
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# img = Image.open('IMG_3579.png')
|
|||
|
|
# img = np.array(img)
|
|||
|
|
|
|||
|
|
# learnWb(img)
|
|||
|
|
|
|||
|
|
# # 读取图像
|
|||
|
|
# image = cv2.imread('IMG_3578.png')
|
|||
|
|
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|||
|
|
#
|
|||
|
|
# # 选择参考色块的区域
|
|||
|
|
# img_patch = image[800:850, 1800:1850, :]
|
|||
|
|
# # 显示图像和参考色块的矩形框
|
|||
|
|
# fig, ax = plt.subplots(figsize=(10, 10))
|
|||
|
|
# ax.set_title('参考块在红色矩形框内')
|
|||
|
|
# ax.imshow(image)
|
|||
|
|
# ax.add_patch(Rectangle((1800, 800), 50, 50, edgecolor='r', facecolor='none'))
|
|||
|
|
#
|
|||
|
|
# # 调用 Ground Truth Algorithm
|
|||
|
|
# ground_truth(image, img_patch, mode='mean')
|
|||
|
|
|
|||
|
|
gray_world("IMG_3577.png")
|