朴实无华的Dataset

朴实无华的Dataset

自定义Dataset，该Dataset用于读取训练图像和标签图像，并做简单的尺寸调整和归一化处理，不包含其他的图像预处理操作。

具体的实例化和使用方法，可以参考main.py的使用。

SimpleDataset.py

import torch
import torch.nn as nn
import math
import torch.nn.functional as F
from PIL import Image
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
import torchvision.transforms as transforms
import cv2
import os

def letterbox_image(image, size):
    iw, ih = image.size
    w, h = size
    scale = min(w/iw, h/ih)
    nw = int(iw*scale)
    nh = int(ih*scale)

    image = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (0,0,0))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image, scale

def transform_small(params=None, grayscale=False, method=Image.BICUBIC, convert=True):
    transform_list = []

    if convert:
        transform_list += [transforms.ToTensor()]
        if grayscale:
            transform_list += [transforms.Normalize((0.5,), (0.5,))]
        else:
            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
    return transforms.Compose(transform_list)

class SimpleDataset(Dataset):
    def __init__(self,img_path, label_path, image_size,num_classes,random_data):
        super(SimpleDataset, self).__init__()

        self.train_lines = os.listdir(img_path)

        self.train_batches = len(self.train_lines)

        self.img_path = img_path
        self.label_path = label_path

        self.image_size = image_size
        self.num_classes = num_classes
        self.random_data = random_data
        self.trans = transform_small()

    def __getitem__(self, index):
        if index == 0:
            shuffle(self.train_lines)
            
        filename = self.train_lines[index]

        # 从文件中读取图像
        img = Image.open(self.img_path + '/' + filename).convert('RGB')
        label = Image.open(self.label_path + '/' + filename).convert('RGB')

        img, scale_A = letterbox_image(img, self.image_size[:2])
        label, scale_B = letterbox_image(label, self.image_size[:2])

        img = self.trans(img)
        label = self.trans(label)

        return img, label, filename

    def __len__(self):
        return self.train_batches

main.py

...
# ===> path
img_path = ''
label_path = ''

# ===> 载入预训练权重
# checkpoint_path = "epoch_8.pth" 
# model_name = "model"  

# logger.info('Loading weights into state dict...')
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# model_dict = eval(model_name).state_dict()
# pretrained_dict = torch.load(checkpoint_path, map_location=device)
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
# model_dict.update(pretrained_dict)
# eval(model_name).load_state_dict(model_dict)
# logger.info('Load finished!')

# ===> 创建优化器和DataLoader
optimizer = optim.Adam(model.parameters(),lr)

train_dataset = SimpleDataset(img_path, label_path, inputs_size, NUM_CLASSES, True)

train_dataloader = DataLoader(train_dataset, batch_size=Batch_size, num_workers=2, pin_memory=True,
                        drop_last=True)

def train(model, dataloader, optimizer, logger, epoch, Epoches, work_dir):
    model = model.train()
    criterion = torch.nn.MSELoss()

    running_loss = 0
    process_loss = 0

    start_time = time.time()
    for iter, batch in enumerate(dataloader):

        imgs, labels, filenames = batch

        imgs = imgs.cuda()
        labels = labels.cuda() 

        optimizer.zero_grad()

        outputs = model(imgs)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()
        running_loss += loss.item() * imgs.size(0)
        process_loss += loss.item() * imgs.size(0)
        if (iter+1)*imgs.size(0) % 50 == 0:
            logger.info('Epoch {}/{} [{}/{}]: loss: {:.4f}'.format(epoch + 1, Epoches, (iter+1)*imgs.size(0), len(dataloader.dataset), process_loss / 100.0))
            process_loss = 0

    epoch_loss = running_loss / len(dataloader.dataset)
    logger.info('Save checkpoint as ' + 'epoch_' + str(epoch) + '.pth...')
    torch.save(model.state_dict(), work_dir + '/' + 'epoch_' + str(epoch) + '.pth')
    logger.info('+------------------------------------------+')
    logger.info('Epoch {}/{}: loss:{:.4f}'.format(epoch + 1, Epoches, epoch_loss))
    logger.info('+------------------------------------------+')