自定义脚本转.raw为16-bit TIFF时偏绿问题求助

我正尝试将.raw格式图像转换为16-bit TIFF格式，但所用的.raw格式无法被rawpy、imageio等库支持，因此使用了自定义预处理脚本。目前转换可完成，但输出图像颜色异常，呈现偏绿状态，无法定位代码中的问题，特此求助。

import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from glob import glob
from tqdm import tqdm

BIT8 = 2 ** 8
BIT16 = 2 ** 16
BIT24 = 2 ** 24

class Debayer3x3(torch.nn.Module):
    def __init__(self):
        super(Debayer3x3, self).__init__()
        # Initialize kernel and index parameters for debayering
        self.kernels = torch.nn.Parameter(
            torch.tensor([
                [0, 0, 0], [0, 1, 0], [0, 0, 0],
                [0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0],
                [0.25, 0, 0.25], [0, 0, 0], [0.25, 0, 0.25],
                [0, 0, 0], [0.5, 0, 0.5], [0, 0, 0],
                [0, 0.5, 0], [0, 0, 0], [0, 0.5, 0],
            ]).view(5, 1, 3, 3), requires_grad=False
        )
        self.index = torch.nn.Parameter(
            torch.tensor([
                [0, 3], [4, 2], [1, 0],
                [0, 1], [2, 4], [3, 0],
            ]).view(1, 3, 2, 2), requires_grad=False
        )

    def forward(self, x):
        # Apply convolution and index gathering for debayering
        B, C, H, W = x.shape
        x = F.pad(x, (1, 1, 1, 1), mode='replicate')
        c = F.conv2d(x, self.kernels, stride=1)
        rgb = torch.gather(c, 1, self.index.repeat(B, 1, H//2, W//2))
        return rgb

def read_raw_24b(file_path, img_shape=(1, 1, 1856, 2880), read_type=np.uint8):
    # Read the raw image file into a suitable format for processing
    raw_data = np.fromfile(file_path, dtype=read_type)
    raw_data = raw_data[0::3] + raw_data[1::3] * BIT8 + raw_data[2::3] * BIT16
    raw_data = raw_data.reshape(img_shape).astype(np.float32)
    return raw_data

def white_balance(image):
    # Calculate the mean of each channel
    mean_r = image[:, :, 0].mean()
    mean_g = image[:, :, 1].mean()
    mean_b = image[:, :, 2].mean()
    
    # Adjust each channel to balance the colors
    image[:, :, 0] *= (mean_g / mean_r)
    image[:, :, 2] *= (mean_g / mean_b)
    
    return image

def func(filename, debayer, out_path, pattern='RGGB', BIT16=65535):
    # Main processing function to read, debayer, and white balance an image
    im = read_raw_24b(filename)
    print(im.shape)
    im = torch.from_numpy(im)

    # Apply the debayer process
    with torch.no_grad():
        im = debayer(im).detach().numpy()

    # Apply white balance
    im = white_balance(im).squeeze()
    im = np.transpose(im, (1, 2, 0))
    im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
    
    # Resize and normalize the image
    print(im.shape)
    im = cv2.resize(im, (720, 720), interpolation=cv2.INTER_LINEAR)
    im = np.clip(im, 0, BIT16 - 1) / (BIT16 - 1)
    im = np.clip(im, 0, 1) * (BIT16 - 1)
    im = im.astype(np.uint16)

    # Save the processed image
    save_path = os.path.join(out_path, os.path.basename(filename).replace('.raw', '.tiff'))
    cv2.imwrite(save_path, im)

def main(in_path, out_path):
    # Set up directories and process each file
    assert os.path.isdir(in_path), f'Invalid input path: {in_path}'
    os.makedirs(out_path, exist_ok=True)
    files = glob(os.path.join(in_path, '*.raw'))
    debayer = Debayer3x3()

    for file_path in tqdm(files):
        func(file_path, debayer, out_path)
        
if __name__ == "__main__":
    in_path = '/home/singh.risha'  # Change to your raw file path
    out_path = '/home/singh.risha'  # Change to your actual output directory path

    main(in_path, out_path)

问题分析与修复方案 #

1. 自定义Debayer模块逻辑错误（核心问题） #

你的Debayer3x3模块输出的图像分辨率是原尺寸的1/2（H//2, W//2），这不符合标准拜耳转换逻辑——标准Debayer应该从单通道拜耳阵列生成同尺寸的三通道RGB图像。这种错误的降采样会导致通道插值混乱，直接引发偏绿问题。

修复：
替换自定义Debayer为OpenCV的原生拜耳转换（更可靠），修改func函数中的Debayer步骤：

def func(filename, out_path, pattern='RGGB', BIT16=65535):
    # 读取RAW并转换为单通道16位数组
    im = read_raw_24b(filename).squeeze().astype(np.uint16)
    # 根据实际拜耳格式选择对应转换flag
    # RGGB对应cv2.COLOR_BAYER_RG2RGB，BGGR对应cv2.COLOR_BAYER_BG2RGB
    rgb = cv2.cvtColor(im, cv2.COLOR_BAYER_RG2RGB)
    
    # 后续白平衡、格式转换步骤不变
    rgb = white_balance(rgb)
    im = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    im = cv2.resize(im, (720, 720), interpolation=cv2.INTER_LINEAR)
    im = np.clip(im, 0, BIT16 - 1).astype(np.uint16)
    
    save_path = os.path.join(out_path, os.path.basename(filename).replace('.raw', '.tiff'))
    cv2.imwrite(save_path, im)

2. RAW数据字节顺序可能错误 #

read_raw_24b中假设字节顺序是低位在前，但很多24位RAW格式是高位字节在前，这会导致数值解析错误，进而影响白平衡计算。

修复：
尝试反转字节顺序：

def read_raw_24b(file_path, img_shape=(1, 1, 1856, 2880), read_type=np.uint8):
    raw_data = np.fromfile(file_path, dtype=read_type)
    # 改为高位字节在前的计算方式
    raw_data = raw_data[2::3] + raw_data[1::3] * BIT8 + raw_data[0::3] * BIT16
    raw_data = raw_data.reshape(img_shape).astype(np.float32)
    return raw_data

3. 白平衡逻辑易受极端值干扰 #

当前用全局均值计算白平衡，若图像存在高亮/暗部极端区域，会导致调整过度。

修复：
改用百分位均值过滤极端值：

def white_balance(image):
    # 取各通道10%-90%区间的像素均值
    g_mean = None
    for c in range(3):
        channel = image[:, :, c].flatten()
        channel = np.sort(channel)
        lower = int(len(channel)*0.1)
        upper = int(len(channel)*0.9)
        # 先计算绿色通道的基准均值
        if c == 1:
            g_mean = channel[lower:upper].mean()
            continue
        # 调整当前通道
        channel_mean = channel[lower:upper].mean()
        image[:, :, c] = image[:, :, c] * (g_mean / channel_mean)
    return image

4. 归一化步骤冗余 #

原代码中两次np.clip和缩放是冗余的，可简化为：

im = np.clip(im, 0, BIT16 - 1).astype(np.uint16)

内容的提问来源于stack exchange，提问作者Hrithik Kanoje