无法为骨折图像应用GridMask，维度广播不匹配错误求助

尝试为骨折图像应用GridMask以缓解数据集不平衡问题，但运行代码时触发数组形状不匹配错误。

原代码 #

import cv2
import os
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator

grid_mask = np.zeros((5,5,1))
grid_mask[2, 2] = 1

dataset_path = "/content/drive/MyDrive/FracAtlas/images/Fractured/"
destination_folder = "/content/drive/MyDrive/FracAtlas/images/masked/masked_images"

os.makedirs(destination_folder, exist_ok=True)

images = []
for filename in os.listdir(dataset_path):
  image = cv2.imread(os.path.join(dataset_path, filename))
  images.append(image)

for i, image in enumerate(images):
  masked_image = image * grid_mask
  cv2.imwrite(os.path.join(destination_folder, f"masked_image_{i}.jpg"), masked_image)

报错信息 #

ValueError                                Traceback (most recent call last)
<ipython-input-20-affc2114a071> in <cell line: 23>()
     22 # Apply the grid mask to each image and save it
     23 for i, image in enumerate(images):
---&gt; 24   masked_image = image * grid_mask
     25   cv2.imwrite(os.path.join(destination_folder, f"masked_image_{i}.jpg"), masked_image)

ValueError: operands could not be broadcast together with shapes (454,373,3) (5,5,1)

错误原因 #

报错核心是数组形状不兼容：

• 输入的骨折图像是彩色图，形状为(高度, 宽度, 3)（示例中是454×373×3）
• 手动定义的grid_mask是固定的5×5×1小尺寸单通道数组
• NumPy的广播规则要求两个数组在每个维度上要么尺寸相同，要么其中一个为1，这里高度、宽度维度完全不匹配，无法执行逐元素相乘。

解决方案 #

需要生成与目标图像尺寸完全匹配的GridMask，并且保证通道数和图像一致（彩色图为3通道）。GridMask的核心是生成周期性的网格遮挡，以下是修正后的代码：

import cv2
import os
import numpy as np

def generate_grid_mask(image_shape, block_size=32, mask_ratio=0.5):
    """生成与输入图像尺寸匹配的GridMask"""
    h, w, c = image_shape
    # 生成单通道网格掩码
    mask = np.ones((h, w, 1), dtype=np.float32)
    
    # 计算网格块的数量
    num_blocks_h = h // block_size + 1
    num_blocks_w = w // block_size + 1
    
    # 交替遮挡网格块
    for i in range(num_blocks_h):
        for j in range(num_blocks_w):
            if (i + j) % 2 == 0:
                # 计算当前块的坐标范围
                y_start = i * block_size
                y_end = min((i+1)*block_size, h)
                x_start = j * block_size
                x_end = min((j+1)*block_size, w)
                mask[y_start:y_end, x_start:x_end] = mask_ratio
    
    # 扩展为3通道，匹配彩色图像
    mask = np.repeat(mask, c, axis=-1)
    return mask

dataset_path = "/content/drive/MyDrive/FracAtlas/images/Fractured/"
destination_folder = "/content/drive/MyDrive/FracAtlas/images/masked/masked_images"

os.makedirs(destination_folder, exist_ok=True)

# 遍历图像并应用GridMask
for idx, filename in enumerate(os.listdir(dataset_path)):
    image_path = os.path.join(dataset_path, filename)
    image = cv2.imread(image_path)
    if image is None:
        print(f"跳过无法读取的图像: {filename}")
        continue
    
    # 生成匹配尺寸的GridMask
    grid_mask = generate_grid_mask(image.shape, block_size=40, mask_ratio=0.0)
    # 应用掩码（逐元素相乘）
    masked_image = image * grid_mask
    # 转换为uint8类型以保存
    masked_image = masked_image.astype(np.uint8)
    
    # 保存处理后的图像
    save_path = os.path.join(destination_folder, f"masked_image_{idx}.jpg")
    cv2.imwrite(save_path, masked_image)

关键修改点 #

• generate_grid_mask函数：根据输入图像的尺寸动态生成掩码，避免形状不匹配问题
• 通道匹配：将单通道掩码扩展为3通道，和彩色图像的通道数对齐
• 周期性网格：按照block_size定义网格块大小，交替遮挡区域，符合GridMask的原始设计逻辑（而非仅中心一个点）
• 增加异常处理：跳过无法读取的图像，避免程序崩溃

内容的提问来源于stack exchange，提问作者Savaş Özcan