U-Net语义分割模型域内iou达0.97，域外新图像测试失效问题排查

我有一个加载了自编码器预训练权重的U-Net模型，该自编码器基于包含1400张图像的数据集训练得到。我使用1400张标注临床数据集开展语义分割任务，模型在自有测试集上iou_score可达0.97，但测试数据集外的随机图像时分割效果极差，不清楚问题成因，代码如下：

数据集与标签训练代码 #

import cv2
import numpy as np
from matplotlib import pyplot as plt
#########################################################################
#Load data for U-net training. 
#################################################################
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ["SM_FRAMEWORK"] = "tf.keras"
import glob
import cv2
import os
import numpy as np
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
train_images = []
#Resizing images is optional, CNNs are ok with large images
SIZE_X = 256 #Resize images (height  = X, width = Y)
SIZE_Y = 256

#Capture training image info as a list

directory_path = '/content/drive/MyDrive/Colab Notebooks/semantic/images/'
list_of_files = sorted( filter( os.path.isfile, glob.glob(directory_path + '*.jpg', recursive=True) ) )

for img_path in list_of_files:
    #for img_path in glob.glob(os.path.join(directory_path, "*.png")):
   print(img_path)
   img = cv2.imread(img_path, cv2.IMREAD_COLOR)               
   img = cv2.resize(img, (SIZE_Y, SIZE_X))
   img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)        
   train_images.append(img)
        
        #train_labels.append(label)
#Convert list to array for machine learning processing        
train_images = np.array(train_images)



train_masks = []
labels_path = '/content/drive/MyDrive/Colab Notebooks/semantic/lables/'
list_of_labels = sorted( filter( os.path.isfile, glob.glob(labels_path + '*.png', recursive=True) ) )
for mask_path in list_of_labels:
    #for img_path in glob.glob(os.path.join(directory_path, "*.png")):
   print(mask_path)
   mask = cv2.imread(mask, 0)               
   mask = cv2.resize(mask, (SIZE_Y, SIZE_X))     
   train_masks.append(mask)     
#Convert list to array for machine learning processing          
train_masks = np.array(train_masks)
#Normalize images
image_dataset = np.array(train_images)/255.
#D not normalize masks, just rescale to 0 to 1.
mask_dataset = np.expand_dims((np.array(train_masks)), 3) /255.
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(image_dataset, mask_dataset, test_size = 0.20, random_state = 0)
#Load unet model and load pretrained weights
from models import build_autoencoder, build_encoder, build_unet
from tensorflow.keras.optimizers import Adam
#import segmentation_models as sm

input_shape = (256, 256, 3)
pre_trained_unet_model = build_unet(input_shape)
pre_trained_unet_model.load_weights('/content/drive/MyDrive/Colab Notebooks/semantic/unet_clinical_model_weights.h5')
pre_trained_unet_model_weights = pre_trained_unet_model.get_weights()[0][1]


pretrained_encoder_wts = np.load('/content/drive/MyDrive/Colab Notebooks/semantic/pretrained_clinical_encoder-weights_300e.npy')

if pre_trained_unet_model_weights.all() == pretrained_encoder_wts.all():
    print("Both weights are identical")
else: 
    print("Something wrong, weghts are different")



pre_trained_unet_model.compile('Adam', loss=sm.losses.binary_focal_jaccard_loss, metrics=[sm.metrics.iou_score])

####################################################################


#Train the model
batch_size=16
pre_trained_unet_model_history = pre_trained_unet_model.fit(X_train, y_train, 
                    verbose=1,
                    batch_size = batch_size,
                    validation_data=(X_test, y_test ), 
                    shuffle=False,
                    epochs=300)

pre_trained_unet_model.save('/content/drive/MyDrive/Colab Notebooks/semantic/pre_trained_unet_model_300epochs.h5')

域内测试代码 #

from keras.models import load_model
pre_trained_unet_model = load_model('/content/drive/MyDrive/Colab Notebooks/semantic/pre_trained_unet_model_300epochs.h5', compile=False)
my_model = pre_trained_unet_model
import random
test_img_number = random.randint(0, X_test.shape[0]-1)
#test_img_number = 119
test_img = X_test[test_img_number]
ground_truth=y_test[test_img_number]
 
test_img_input=np.expand_dims(test_img, 0)
prediction = (my_model.predict(test_img_input)[0,:,:,0] > 0.5).astype(np.uint8)
 
plt.figure(figsize=(16, 8))
plt.subplot(231)
plt.title('Testing Image')
plt.imshow(test_img, cmap='gray')
plt.subplot(232)
plt.title('Testing Label')
plt.imshow(ground_truth[:,:,0], cmap='gray')
plt.subplot(233)
plt.title('Prediction on test image')
plt.imshow(prediction, cmap='gray') 
plt.show()

域内测试效果：

域外测试代码 #

from keras.models import load_model
import segmentation_models as sm
import numpy as np
import matplotlib.pyplot as plt

model = load_model('/content/drive/MyDrive/Colab Notebooks/semantic/pre_trained_unet_model_300epochs.h5', compile=False) # load the model
model.compile(loss=sm.losses.binary_focal_jaccard_loss, optimizer='Adam', metrics=[sm.metrics.iou_score])
from keras.preprocessing import image

test_image= image.load_img('/content/drive/MyDrive/Colab Notebooks/semantic/images/Foot ulcer 3-3.jpg',target_size = (256, 256))

test_image = image.img_to_array(test_image)
test_image = np.expand_dims(test_image, axis = 0)
result = model.predict(test_image)
result_img = result.reshape(256,256)
plt.imshow(result_img, cmap='gray')

域外测试效果：

一、测试流程错误（直接导致域外预测失效） #

• 预处理逻辑不一致：训练阶段使用cv2.imread读取图像，默认是BGR通道顺序，后续做了cv2.COLOR_RGB2BGR转换，同时将像素值归一化到[0,1]区间；但域外测试使用keras.preprocessing.image.load_img读取图像，默认是RGB通道顺序，既没有做通道转换，也没有执行除以255的归一化操作，输入数据分布和训练时完全不匹配，自然预测效果极差。
• 预测后处理缺失：域内测试时对预测结果做了>0.5的阈值二值化，域外测试直接输出原始概率图，也是可视化效果差的原因之一。

二、模型泛化性缺陷（导致域外真实数据效果不佳） #

• 数据集分布偏差：自有数据集的1400张图像是同分布采集，随机拆分的训练/测试集分布完全一致，所以域内测试IOU很高，但域外数据在光照、拍摄角度、病灶形态、背景、拍摄设备等维度存在差异，模型没有学习到通用的溃疡特征。
• 训练设置不合理：训练时fit函数的shuffle参数设为False，模型会按固定顺序学习样本，容易过拟合域内数据的特定特征，无法泛化到新数据。
• 缺少数据增强：整个训练流程没有加入随机翻转、旋转、亮度调整、噪声注入等数据增强操作，模型对样本变化的鲁棒性极差。
• 预训练没有提升泛化性：自编码器预训练使用的数据集和分割任务的数据集完全相同，相当于只在域内数据上做了特征预训练，无法提供跨域的特征表达能力。

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