fastai_UNet.py

from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from sklearn.model_selection import train_test_split
# from fastai.vision.models import resnet34
# from fastai.vision.leaner import unet_learner
from fastai.vision import *
from torch.utils.data import Dataset
from torchvision import transforms, utils



df = pd.read_csv('/path/to/csv/data.csv')

X = list(df['input_img'])
y = list(df['mask_img'])

X_train, X_valid, y_train, y_valid = train_test_split(
     X, y, test_size=0.33, random_state=42)



class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, img):

        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        img = img.transpose((2, 0, 1))
        return {'image': torch.from_numpy(img),
                }


class NumbersDataset(Dataset):
    def __init__(self, inputs, labels, transform=None):
        classes = [0,1]
        self.X = inputs
        self.y = labels
        self.transform = transform
        self.c = 2

    def __len__(self):
        return len(self.X)

   def __getitem__(self, idx):
        img_train = cv2.imread(self.X[idx])
        img_mask = cv2.imread(self.y[idx])
        img_train = cv2.resize(img_train, (427,240), interpolation = cv2.INTER_LANCZOS4) 
        img_mask = cv2.resize(img_mask, (427,240), interpolation = cv2.INTER_LANCZOS4)
        img_mask = cv2.cvtColor(img_mask, cv2.COLOR_BGR2GRAY)
        bin_mask = np.zeros_like(img_mask)
        bin_mask[(img_mask)>0]=1
        bin_mask = bin_mask.reshape(240, 427, 1)
        # print("################################################################")
        # return self.X[idx], self.y[idx]
        if self.transform:
            img_train = self.transform(img_train)
            bin_mask = self.transform(bin_mask)

        return img_train, bin_mask


if __name__ == '__main__':
    # dataset_train = NumbersDataset(X_train, y_train, transforms.Compose([ToTensor()]))
    dataset_train = NumbersDataset(X_train, y_train)
    dataloader_train = DataLoader(dataset_train, batch_size=4, shuffle=True)

    dataset_valid = NumbersDataset(X_valid, y_valid)
    # dataset_valid = NumbersDataset(X_valid, y_valid, transforms.Compose([ToTensor()]))
    dataloader_valid = DataLoader(dataset_valid, batch_size=4, shuffle=True)

#     datas = DataBunch.create(train_ds = dataloader_train, valid_ds = dataloader_valid)
    datas = DataBunch(train_dl = dataloader_train, valid_dl = dataloader_valid)
#     datas.show_batch()
    # res = datas.c
    datas.c = 1
    learner = unet_learner(datas, models.resnet34)
    # print(len(dataset))
    # plt.imshow(dataset[100])
    # plt.show()
    # print(next(iter(dataloader)))
    # model = resnet34()
    # print(model)
    # for i, batch in enumerate(dataloader_train):
        # print(i, batch)
        # inp ,label = batch
        # print(inp.shape, label.shape)  

        # leaner = unet_learner(data = batch, arch = models.resnet34)