使用DINOv2生成图像分割蒙版后保存全黑的问题求助

问题描述 #

使用DINOv2模型开展图像分割任务，基于patch embeddings生成蒙版。可视化时蒙版能正常显示，但保存为PNG或JPG格式后全部变为黑色，尝试转为RGB等方法仍无法解决。

相关代码 #

import numpy as np
import torch.nn.functional as F
ans=[]
masks=[]
for idx,(image,_) in enumerate(trainloader):
    seg=segmentor(1)
    x=seg.forward(layerpatches)
    print(f"mask image shape:{x.shape}")
    for i in range(x.shape[0]):
        image_array = x[i].detach().squeeze().numpy()
        mask_image = image_array
        image=image.permute(0,2,3,1).squeeze()
        fig, ax = plt.subplots(1, 2, figsize=(10, 5))
        ax[0].imshow(mask_image)
        ax[0].set_title('Mask')
        ax[0].axis('off')
        ax[1].imshow(image)
        ax[1].set_title('Corresponding Image')
        ax[1].axis('off')
        plt.show()
        mask_name = f"{img_name[idx]}_mask.jpg"
        mask_np = np.clip(mask_rgb, 0, 255).astype(np.uint8)
        # # Save the adjusted mask as PNG
        mask_path = os.path.join(mask_dir, mask_name)
        cv2.imwrite(mask_path,mask_image)
    break

现象对比 #

• 可视化状态：蒙版与对应图像显示正常，蒙版细节清晰可见
• 保存后状态：蒙版文件全部显示为黑色，无任何有效内容

问题原因及解决方法 #

核心原因 #

plt.imshow会自动将蒙版的数值范围归一化到0-1区间进行显示，但cv2.imwrite要求输入的图像数据必须是0-255范围的uint8类型。代码中直接保存模型输出的原始数组（通常是浮点数，范围并非0-255），导致保存后无法正常显示。此外代码中mask_rgb未定义，且处理后的mask_np并未实际用于保存，逻辑存在无效环节。

修改方案 #

对蒙版数组进行数值归一化和类型转换后再保存，替换原代码中的保存逻辑：

# 替换原保存代码段
mask_name = f"{img_name[idx]}_mask.jpg"
# 1. 归一化蒙版数值到0-1范围（适配任意范围的模型输出）
mask_normalized = (mask_image - mask_image.min()) / (mask_image.max() - mask_image.min())
# 2. 转换为0-255的uint8类型
mask_uint8 = (mask_normalized * 255).astype(np.uint8)
# 3. 保存蒙版
mask_path = os.path.join(mask_dir, mask_name)
cv2.imwrite(mask_path, mask_uint8)

如果蒙版是二值分割结果（仅区分前景/背景），可直接用阈值处理：

# 假设阈值为0.5，根据模型输出实际情况调整
mask_binary = (mask_image > 0.5).astype(np.uint8) * 255
cv2.imwrite(mask_path, mask_binary)

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