【Kaggle代码】Carvana Image Masking Challenge

Kaggle比赛链接：https://www.kaggle.com/competitions/carvana-image-masking-challenge

提交的时候老是0分，最终发现是图像预处理的时候将其大小调整为了256*256，但是其实际大小是1280,*1918，在训练集输出结果的时候将图像还原即可。

积累的常用代码

• pytorch中，输入模型的数据必须是tensor类型。所以在Dataloader中，最后要将其转换为tensor类型。

• self.images = sorted(os.listdir(image_dir))

  os.listdir返回的是一个列表，sorted是将列表中按照文件名进行排序。

•  gif_reader = imageio.get_reader(mask_path)
   single_frame = gif_reader.get_data(0)`
  

  用于读取gif图片，读取的是其第一帧。此时的singe_frame已经是图片的形式了，而非gif形式。

• train_dataset,valid_dataset = random_split(dataset, [0.8,0.2])

  random_split 函数是 PyTorch 库中的 torch.utils.data 模块的一部分，用于将数据集随机划分成给定的子集。导入语句：:from torch.utils.data import random_split

• train_features_batch, train_labels_batch = next(iter(train_dataloader))

  next(iter(train_dataloader)) 相当于是迭代一次，一次获取的图像数量是一个 batch_size。

• 最终的提交是用了RLE编码来压缩图像 。举个简单的例子：这里有一个一位图列表表示为：one=[1,1,1,1,0,0,0,0,0,1,1,1,1,1,1]，每个元素为一像素。他可以写成：(4,1), (5,0), (6,1)，表示有分别有4个1，5个0，6个1连续出现。我们把4，5，6变成二进制格式然后算一下编码后省下了多少空间…省下三比特位。参考博客：行程编码（RLE）

代码

代码如下，按照 ipynb结构排列：

import numpy as np
import os
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
import matplotlib.pyplot as plt
from torch import optim
from torch import nn
import cv2 
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch.utils.data import random_split
from tqdm import tqdm
import imageio
import torchvision
import matplotlib.pyplot as plt
from PIL import Image # Image模块是在Python PIL图像处理中常见的模块，对图像进行基础操作的功能基本都包含于此模块内。

import zipfile
# Unzip Files
with zipfile.ZipFile('/kaggle/input/carvana-image-masking-challenge/train.zip', 'r') as zip_ref:
    zip_ref.extractall('/kaggle/working/')

with zipfile.ZipFile('/kaggle/input/carvana-image-masking-challenge/train_masks.zip', 'r') as zip_ref:
    zip_ref.extractall('/kaggle/working/')
with zipfile.ZipFile('/kaggle/input/carvana-image-masking-challenge/test.zip', 'r') as zip_ref:
    zip_ref.extractall('/kaggle/working/')

os.listdir('/kaggle/working')

image_dir = '/kaggle/working/train'
mask_dir = '/kaggle/working/train_masks'
len(os.listdir(image_dir)), len(os.listdir(mask_dir))

class MyDataset(Dataset):
    def __init__(self,image_dir,mask_dir,width=320,high=380,transforms=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.width = width
        self.high = high
        self.transforms = transforms
        
        
        self.images = sorted(os.listdir(image_dir))
        self.masks = sorted(os.listdir(mask_dir))
        #sorted() 函数对这些文件名进行排序，以确保它们按字母顺序排列。
        #这对于在数据集中保持一致的顺序很重要，以便图像和对应的掩码匹配。
        #最终得到的images是具体图片的名称，用于后面图像绝对地址的拼接。mask同理
        
    def __len__(self):
        return len(self.images)
        
    def __getitem__(self,idx):
        img_path = os.path.join(self.image_dir, self.images[idx])  # 将image_dir与图像列表中的每张图的名字连接成地址
        mask_path = os.path.join(self.mask_dir, self.masks[idx])  # 将mask_dir与标签列表中的每张图的名字连接成地址
        
         #转化为灰度图
        image_data = cv2.imread(img_path)
        image_data = cv2.cvtColor(image_data, cv2.COLOR_BGR2GRAY)
        
        gif_reader = imageio.get_reader(mask_path)
        single_frame = gif_reader.get_data(0)
        mask_data = cv2.cvtColor(single_frame, cv2.COLOR_BGR2GRAY)
        
        #图像尺寸调整
        dimensions = (self.width, self.high)
        image_data = cv2.resize(image_data, dimensions, interpolation=cv2.INTER_AREA)
        mask_data = cv2.resize(mask_data, dimensions, interpolation=cv2.INTER_AREA)
        
        #转化为tensor类型，ToTensor会自动将其值缩放到0-1之间
        image_data = transforms.ToTensor()(image_data)
        mask_data = transforms.ToTensor()(mask_data)
        
        
        
        
#         #转为numpy数组类型更便于对数据的处理
#         image = np.array(Image.open(img_path).convert('RGB'))
#         mask = np.array(Image.open(mask_path).convert('L'), dtype=np.float32)
        
#         # 对掩码数组进行了处理，将所有值为255.0的像素点（假设是表示目标的像素）设置为1.0。
#         #这样，掩码数组中的像素值就被转换为了二进制格式，其中1表示目标区域，0表示背景区域。
#         mask[mask == 255.0] = 1.0

#         #这里得mask本来就是0-1二值图，如果不是要转为0-1二值图得话，用下面得代码更合理、
#         #  mask[mask >= 1.0] = 1.0

        
#         if self.transform is not None:

#             augmentations = self.transform(image=image, mask=mask)
#             #这里是集中一次性对image和mask都进行了transforms转换，最终得到的augmentations是一个字典。
#             #参数前一个image是字典的关键字，后一个image是传入的图像数据。
#             #字典类型是：'image': <增强后的图像数据>,'mask': <增强后的掩码数据>  
            
#             #这里再通过关键字读取字典中的值
#             image = augmentations['image']
#             mask = augmentations['mask']

#             #上述操作和以下代码操作结果一样，只是上面代码的一次性进行了image和mask的操作
# #             image = transform(image)
# #             mask = transform(mask)

        
        
        return image_data,mask_data

#划分数据集，将原来的训练集划分为  训练集和验证集
dataset = MyDataset(image_dir,mask_dir,256,256)
train_dataset,valid_dataset = random_split(dataset, [0.8,0.2])

#创建数据加载器
train_dataloader = DataLoader(train_dataset,batch_size=4,shuffle=True,num_workers=8)
valid_dataloader = DataLoader(valid_dataset,batch_size=4,shuffle=True,num_workers=8)
print(len(train_dataset))
print(len(valid_dataset))

#可视化一个batch的图像


train_features_batch, train_labels_batch = next(iter(valid_dataloader))
#next(iter(train_dataloader)) 相当于是迭代一次，一次获取的图像数量是一个batch_size


print(train_features_batch.shape, train_labels_batch.shape)

#显示单张图片
plt.subplot(1,2,1)
plt.imshow(train_features_batch[0,0], cmap='gray')
plt.subplot(1,2,2)
plt.imshow(train_labels_batch[0,0], cmap='gray')
#torch.Size([4, 1, 256, 256]) ，因为是四维的，所有要指定出前两维train_features_batch[0,0]


# 创建一个图像窗口,用于显示4张图片
plt.figure(figsize=(12, 4))  # 调整图像窗口的大小

# 循环遍历每个样本并显示
for i in range(4):
    plt.subplot(2, 4, i+1)
    plt.imshow(train_features_batch[i,0], cmap='gray')  # 假设是灰度图
    plt.axis('off')  # 关闭坐标轴
    plt.subplot(2, 4, i+5)
    plt.imshow(train_labels_batch[i,0], cmap='gray')  # 假设是灰度图
   
    plt.axis('off')  # 关闭坐标轴
plt.subplots_adjust(wspace=0.1,hspace=0.1)
plt.show()

#网络搭建
class conv_block(nn.Module):  #定义卷积模块，包含两次卷积
    def __init__(self, ch_in, ch_out):
        super(conv_block, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True),

            # 第二次卷积不改变通道数，但是尺寸会变。
            nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class up_conv(nn.Module):    #上卷积模块  包括一次上采样 一次卷积
    def __init__(self, ch_in, ch_out):
        super(up_conv, self).__init__()
        self.up = nn.Sequential(
            # 不改变通道个数，但是尺寸变为2倍
            nn.Upsample(scale_factor=2),
            # 进行一次卷积
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x

class U_Net(nn.Module):
    def __init__(self, img_ch=1, output_ch=1):
        super(U_Net, self).__init__()

        self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(ch_in=img_ch, ch_out=8)
        self.Conv2 = conv_block(ch_in=8, ch_out=16)
        self.Conv3 = conv_block(ch_in=16, ch_out=32)
        self.Conv4 = conv_block(ch_in=32, ch_out=64)
        self.Conv5 = conv_block(ch_in=64, ch_out=128)

        self.Up5 = up_conv(ch_in=128, ch_out=64)
        self.Up_conv5 = conv_block(ch_in=128, ch_out=64)

        self.Up4 = up_conv(ch_in=64, ch_out=32)
        self.Up_conv4 = conv_block(ch_in=64, ch_out=32)

        self.Up3 = up_conv(ch_in=32, ch_out=16)
        self.Up_conv3 = conv_block(ch_in=32, ch_out=16)

        self.Up2 = up_conv(ch_in=16, ch_out=8)
        self.Up_conv2 = conv_block(ch_in=16, ch_out=8)

        self.Conv_1x1 = nn.Conv2d(8, output_ch, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        # encoding path
        x1 = self.Conv1(x)

        x2 = self.Maxpool(x1)
        x2 = self.Conv2(x2)

        x3 = self.Maxpool(x2)
        x3 = self.Conv3(x3)

        x4 = self.Maxpool(x3)
        x4 = self.Conv4(x4)

        x5 = self.Maxpool(x4)
        x5 = self.Conv5(x5)

        # decoding + concat path
        d5 = self.Up5(x5)
        d5 = torch.cat((x4, d5), dim=1)

        d5 = self.Up_conv5(d5)

        d4 = self.Up4(d5)
        d4 = torch.cat((x3, d4), dim=1)
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        d3 = torch.cat((x2, d3), dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        d2 = torch.cat((x1, d2), dim=1)
        d2 = self.Up_conv2(d2)

        d1 = self.Conv_1x1(d2)
        d1 = torch.sigmoid(d1)

        return d1

def dice_score(preds, targets):
    preds = F.sigmoid(preds)
    preds = (preds > 0.5).float()
    score = (2. * (preds * targets).sum()) / (preds + targets).sum()
    return torch.mean(score).item()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = U_Net(img_ch=1, output_ch=1).to(device)

loss_fn = torch.nn.BCELoss()
optimizer = torch.optim.Adam(params=model.parameters(),lr=1e-3)

#测试一下网络输出是否正确
tensor = torch.rand(4, 1, 256, 256).to(device)  # 将输入数据移动到 GPU 上,torch模型默认输入是tensor B*C*W*H
print(tensor.shape)
out = model(tensor)

print(out.shape)

epochs=10
train_loss_all = [] #定义一个列表用于保存总的训练集loss，方便后续打印
val_loss_all =[]   #定义一个列表用于保存总的验证集loss，方便后续打印
best_loss = 1e10   #记录最佳的loss
for epoch in range(epochs):
    train_loss = 0
    val_loss = 0
    train_num = 0  #总共训练集的图片数量
    val_num = 0  #总共验证集的图片数量
    ##训练集
    model.train() ## 设置模型为训练模式
    for batch, (image, mask) in tqdm(enumerate(train_dataloader)):
       # print(image.shape)
        
        
        optimizer.zero_grad()  ##每次优化前先进行梯度清零
        image = image.to(device) #放到gpu上
        mask = mask.to(device) #放到gpu上
        
        output = model(image)
        loss = loss_fn(output,mask)
        loss.backward()
        optimizer.step() #梯度更新
        
        
        train_loss += loss.item()
        #loss.item()给出当前批次的损失值，len(image)即批次中样本的数量,整个batchsize（4张图片）的总loss
        #这样得到的是当前批次的总损失。这个总损失通过epoch_loss += loss.item()*len(image)  进行累加，
        #最终得到一个epoch中所有批次的训练损失的总和。
        
        train_num += 1
    #计算一个epoch的训练集loss
    train_loss_all.append(train_loss / train_num)
    print('{} Train Loss:{:.4f}'.format(epoch,train_loss_all[-1]))
    
    
    ##验证集
    model.eval() ## 设置模型为训练模式
    for batch, (image, mask) in tqdm(enumerate(valid_dataloader)):
        image = image.to(device) #放到gpu上
        mask = mask.to(device) #放到gpu上
        output = model(image)
        loss = loss_fn(output,mask)
        val_loss += loss.item() 
        val_num += 1
     #计算一个epoch的训练集loss
    val_loss_all.append(val_loss / val_num)
    print('{} Valid Loss:{:.4f}'.format(epoch,val_loss_all[-1]))
    
    ##保存模型
    if val_loss_all[-1] < best_loss :
        best_loss = val_loss_all[-1]
        check_points = model.state_dict()
        torch.save(check_points, '/kaggle/working/BestSave.pt')
    if epoch % 10 == 0:
        check_points = model.state_dict()
        model_name = 'epoch_number%d' % (epoch) + '.pt'
        torch.save(check_points, '/kaggle/working/'+model_name)

#可视化模型训练过程中的loss曲线
epochs = list(range(1, 11))  # 或者任何你实际的 epochs 数量

plt.figure(figsize=(10,6))
plt.plot(epochs,train_loss_all,"ro-",label = "Train Loss")
plt.plot(epochs,val_loss_all,"bs-",label = "Valid Loss")
plt.legend()
plt.xlabel("epoch")
plt.ylabel("Loss")
plt.show()

#使用单张图像查看一下模型训练效果
# 1. 加载训练好的模型
model = model = U_Net(img_ch=1, output_ch=1).to(device)
model.load_state_dict(torch.load('/kaggle/working/BestSave.pt'))
model.eval()

# 2. 预处理单张图像
transform = transforms.Compose([
    
    transforms.Grayscale(),  # 转为灰度图
    transforms.Resize((256, 256)),  # 调整图像尺寸
    transforms.ToTensor(),  # 转为 PyTorch 的 Tensor 格式，同时将像素值缩放到 [0, 1]
])

image_path = '/kaggle/input/carvana-image-masking-challenge/29bb3ece3180_11.jpg'
input_image = Image.open(image_path)
input_tensor = transform(input_image).unsqueeze(0).to(device)  # 添加 batch 维度

# 3. 将输入图像传递给模型进行推理
with torch.no_grad():
    output = model(input_tensor)
#4. 结果可视化
torchvision.utils.save_image((output > 0.0).float(), '/kaggle/working/result1.jpg')
torchvision.utils.save_image((output > 0.5).float(), '/kaggle/working/result2.jpg')



import torchvision.transforms.functional as F
output = (output > 0.5).float().cpu().numpy().squeeze()

# 可视化原始图像和二值化的模型输出
plt.figure(figsize=(8, 4))

# 原始图像
plt.subplot(1, 2, 1)
plt.imshow(F.to_pil_image(input_tensor.squeeze().cpu()))
plt.title('Original Image')

# 二值化的模型输出
plt.subplot(1, 2, 2)
plt.imshow(output, cmap='gray')
plt.title('Binary Output')

plt.show()

#前面的训练集中是包含了mask，但是测试集中没有包含mask，所以dataset会有一些区别。应该删除mask的操作
class CarvanaTestDataset(Dataset):
    def __init__(self,image_dir, width=256,high=256):
        self.image_dir = image_dir
        self.width = width
        self.high = high
        
        self.images = sorted(os.listdir(image_dir))
        #sorted() 函数对这些文件名进行排序，以确保它们按字母顺序排列。
        #这对于在数据集中保持一致的顺序很重要，以便图像和对应的掩码匹配。
        #最终得到的images是具体图片的名称，用于后面图像绝对地址的拼接。mask同理
        
    def __len__(self):
        return len(self.images)
        
    def __getitem__(self,idx):
        img_path = os.path.join(self.image_dir, self.images[idx])  # 将image_dir与图像列表中的每张图的名字连接成地址
        img_name = self.images[idx]
        
        #转化为灰度图
        image_data = cv2.imread(img_path)
        image_data = cv2.cvtColor(image_data, cv2.COLOR_BGR2GRAY)
        
        
        #图像尺寸调整
        dimensions = (self.width, self.high)
        image_data = cv2.resize(image_data, dimensions, interpolation=cv2.INTER_AREA)
       
        
        #转化为tensor类型，ToTensor会自动将其值缩放到0-1之间
        image_data = transforms.ToTensor()(image_data)

        
        return img_name,image_data

test_dir = '/kaggle/working/test'

len(os.listdir(test_dir))

test_dataset = CarvanaTestDataset(test_dir,256,256)
test_dataloader = DataLoader(test_dataset,batch_size=16,shuffle=False,num_workers=4)

import numpy as np
import pandas as pd
import time

def rle_encode(img):
    '''
    img: numpy array, 1 - mask, 0 - background
    Returns run length as string formated
    '''
     # 将二维图像数组展平成一维
    pixels = img.flatten()
     # 将背景像素标记为0
    pixels[0] = 0
    pixels[-1] = 0
    
     # 找到非连续变化的位置
#     pixels[1:]: 这是一个切片，包含了数组中从索引1到最后的所有元素。
#     pixels[:-1]: 这也是一个切片，包含了数组中从开头到倒数第二个元素的所有元素。
#     通过比较这两个切片的元素是否不相等，我们可以找到从一个像素值到下一个像素值发生变化的位置。
#如果两个相邻的像素值相等，pixels[1:] != pixels[:-1] 的结果就是 False，表示没有发生变化；
#如果不相等，结果就是 True，表示发生了变化。

    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
    #  + 2: 由于运行长度编码通常从 1 开始，而不是从 0 开始，因此将所有的索引值加 2，将索引调整为 1-based。
    #  通过 + 2 的操作将索引调整为 1-based，这是因为运行长度编码通常从 1 开始记录位置。
    # 例如，第一个像素的位置通常被表示为 1，而不是 0。
    
    
     # 计算 RLE 编码
    runs[1::2] -= runs[:-1:2]
    #runs[1::2]: 这是一个切片操作，选择 runs 数组中从索引1开始，每隔2个元素的所有元素。这包含了所有变化位置的奇数索引，用于存储运行长度编码的长度。

#   runs[:-1:2]: 这也是一个切片操作，选择 runs 数组中从开头到倒数第二个元素，每隔2个元素的所有元素。这包含了所有变化位置的偶数索引，用于存储运行长度编码的起始位置。

#   runs[1::2] -= runs[:-1:2]: 这是一个减法操作，将每个奇数索引的值减去对应的偶数索引的值。这是因为 RLE 编码中，奇数索引存储的是长度，而偶数索引存储的是起始位置。通过相减，得到最终的运行长度编码。
    
    
    
    
    
     # 将结果以字符串形式返回
    return ' '.join(str(x) for x in runs)

model  = U_Net(img_ch=1, output_ch=1).to(device)
model.load_state_dict(torch.load('/kaggle/working/BestSave.pt'))
model.eval()

# Predictions
all_predictions = []

# loop = tqdm(tem_dataloader)   测试
loop = tqdm(test_dataloader)
for i, (img_names , image) in enumerate(loop):  
    image = image.to(device)
    with torch.no_grad():
        preds = model(image)
        preds = (preds > 0.5).float()
        preds = TF.resize(
        # # 调整预测的大小为指定的大小（在此处为1280x1918），使用最近邻插值
        preds, size=(1280, 1918), interpolation=TF.InterpolationMode.NEAREST
    )
    
     # Encoding
    for idx in range(len(image)):
        encoding = rle_encode(preds[idx].squeeze().cpu())
        all_predictions.append([img_names[idx], encoding])
        
        

import shutil
WORKING_DIR='/kaggle/working/'
shutil.rmtree(WORKING_DIR + 'test')
shutil.rmtree(WORKING_DIR + 'train')
shutil.rmtree(WORKING_DIR + 'train_masks')

import shutil
import pandas as pd


sub = pd.DataFrame(all_predictions)
sub.columns = ['img', 'rle_mask']
sub.to_csv(os.path.join(WORKING_DIR, 'submission.csv'), index=False)

import os
from IPython.display import FileLink

os.chdir('/kaggle/working')
FileLink('submission.csv')