U-Net模型仅输出单类别有效分类报告问题求助

问题现象 #

我的U-Net模型输出的分类报告仅显示单类别有效数据，其余类别support为0，具体报告如下：

Classification Report:
precision recall f1-score support

discbulge 1.00 0.96 0.98 5898240
herniation 0.00 0.00 0.00 0
normal 0.00 0.00 0.00 0

accuracy 0.96 5898240
macro avg 0.33 0.32 0.33 5898240
weighted avg 1.00 0.96 0.98 5898240

附上我的模型代码：

import os
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, Concatenate
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
import matplotlib.pyplot as plt

# Define U-Net architecture
def unet():
    inputs = Input(shape=(height, width, channels))

    # Encoder
    conv1 = Conv2D(64, (3, 3), padding='same', activation='relu')(inputs)
    conv1 = Conv2D(64, (3, 3), padding='same', activation='relu')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(128, (3, 3), padding='same', activation='relu')(pool1)
    conv2 = Conv2D(128, (3, 3), padding='same', activation='relu')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    # Decoder
    up1 = UpSampling2D(size=(2, 2))(pool2)
    up1 = Conv2D(64, (3, 3), padding='same', activation='relu')(up1)
    up1 = Conv2D(64, (3, 3), padding='same', activation='relu')(up1)
    concat1 = Concatenate(axis=-1)([conv2, up1])

    up2 = UpSampling2D(size=(2, 2))(concat1)
    up2 = Conv2D(64, (3, 3), padding='same', activation='relu')(up2)
    up2 = Conv2D(64, (3, 3), padding='same', activation='relu')(up2)
    concat2 = Concatenate(axis=-1)([conv1, up2])

    outputs = Conv2D(num_classes, (1, 1), padding='valid', activation='softmax')(concat2)

    model = Model(inputs=inputs, outputs=outputs)
    return model

# Set up dataset paths and parameters
image_dir = '/content/drive/MyDrive/opp/lumbar/Images'
label_dir = '/content/drive/MyDrive/opp/lumbar/Masks'
#output_dir = '/content/drive/MyDrive/files/Segmented_Images'
num_classes = 3  # Number of segmentation classes (IVD, PE, TS, AAP)
height = 256  # Image height
width = 256  # Image width
channels = 3  # Number of image channels

# Define functions to load and preprocess images and labels
def load_and_preprocess_image(image_path):
    img = Image.open(image_path)
    img = img.resize((width, height))
    img = np.array(img) / 255.0  # Normalize image
    return img

def load_and_preprocess_label(label_path):
    img = Image.open(label_path)
    img = img.resize((width, height))
    img = img.convert('L')  # Convert to grayscale
    img = np.array(img)
    # Perform label encoding
    encoded_labels = np.zeros((height, width, num_classes))
    for c in range(num_classes):
        encoded_labels[:, :, c] = (img == c).astype(int)
    return encoded_labels

# Load and preprocess the dataset
image_filenames = os.listdir(image_dir)
label_filenames = os.listdir(label_dir)

images = []
labels = []

for img_file, lbl_file in zip(image_filenames, label_filenames):
    img_path = os.path.join(image_dir, img_file)
    lbl_path = os.path.join(label_dir, lbl_file)

    # Load and preprocess the image
    image = load_and_preprocess_image(img_path)
    images.append(image)

    # Load and preprocess the label
    label = load_and_preprocess_label(lbl_path)
    labels.append(label)

images = np.array(images)
labels = np.array(labels)

# Split dataset into training and validation sets
train_images, val_images, train_labels, val_labels = train_test_split(images, labels, test_size=0.2)

# Build and compile the U-Net model
model = unet()
model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])

# Train the model
history = model.fit(train_images, train_labels, batch_size=16, epochs=10, validation_data=(val_images, val_labels))

问题根源分析 #

1. 标签映射错误：代码默认标签掩码的灰度值等于类别索引（img == c），但实际掩码文件的灰度值可能和定义的类别索引不匹配，导致herniation、normal类未被正确编码，support为0。
2. 数据集划分不合理：随机划分数据集可能导致验证集中完全缺失部分类别的像素，或者原始数据集本身就存在类别失衡/缺失。
3. 标签加载逻辑缺陷：强制转换为灰度图（convert('L')）可能改变原始掩码的像素值，破坏了类别与像素值的对应关系。

解决步骤 #

1. 验证标签映射关系 #

随机抽取掩码文件，查看实际像素值：

sample_label = Image.open(os.path.join(label_dir, label_filenames[0]))
print(np.unique(sample_label))

根据打印出的唯一值，修正load_and_preprocess_label中的匹配逻辑。比如如果herniation对应灰度值是100，就改成：

encoded_labels[:, :, 1] = (img == 100).astype(int)

2. 检查并修正数据集分布 #

统计全数据集的类别像素占比：

total_pixels = height * width * len(labels)
class_counts = []
for c in range(num_classes):
    count = np.sum(labels[:, :, :, c])
    class_counts.append(count)
    print(f"类别{c}像素数量: {count}, 占比: {count/total_pixels:.4f}")

• 若某类像素为0：需补充对应类别的数据。
• 若仅验证集缺失：改用分层划分，确保验证集包含所有类别：

# 生成每个图像的类别存在标记
image_class_markers = []
for label in labels:
    present_classes = np.any(label, axis=(0,1))
    image_class_markers.append(tuple(present_classes))

# 按标记分层划分数据集
train_images, val_images, train_labels, val_labels = train_test_split(
    images, labels, test_size=0.2, stratify=image_class_markers
)

3. 修正标签加载逻辑 #

避免不必要的灰度转换，直接读取原始掩码像素值：

def load_and_preprocess_label(label_path):
    img = Image.open(label_path)
    img = img.resize((width, height))
    img = np.array(img)
    # 自定义类别-灰度值映射，替换为实际值
    class_mapping = {0: 0, 1: 100, 2: 200}
    encoded_labels = np.zeros((height, width, num_classes))
    for idx, gray_val in class_mapping.items():
        encoded_labels[:, :, idx] = (img == gray_val).astype(int)
    return encoded_labels

4. 调整训练策略应对类别失衡 #

使用加权交叉熵损失，给少数类更高权重：

class_weights = total_pixels / np.array(class_counts)
class_weights = class_weights / np.sum(class_weights)  # 归一化
model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'], class_weight=class_weights)

同时可增加训练轮次，或对包含少数类的图像进行翻转、旋转等数据增强操作。

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