main.py

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torchvision as tv
import torch.optim as optim
import scipy.io
import os
from PIL import Image
import numpy


#neural network
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        #define the properties of each layer
        self.conv = nn.Conv2d(1,1,5, padding=2)

    def forward(self,x):
        #define the forward pass of the nn
        #back propogation gradients will be automatically calculated
        x = self.conv(x)
        return x

#data set images need to be rotated to portrait
def imrotate(x):
    if x.size != (481, 321):
        x = x.rotate(90, expand=True)
    return x;

#path to input images
image_path = './BSR_bsds500/BSR/BSDS500/data/images/train'
truth_path = './BSR_bsds500/BSR/BSDS500/data/groundTruth/train'

#pretransform dataset images while loading
BSR_transform = tv.transforms.Compose([
    tv.transforms.Grayscale(),
    tv.transforms.Lambda(lambda x : imrotate(x)),
    tv.transforms.ToTensor(),
    tv.transforms.Normalize((0.5,),(0.5,))
])

#load input images from path and apply pertransforms
train_set = tv.datasets.ImageFolder(root = image_path, transform = BSR_transform)

ground_truth = []

#load ground truths
for filename in sorted(os.listdir(truth_path)):
    #extract from .mat files
    mat = scipy.io.loadmat(os.path.join(truth_path, filename))
    mat = mat['groundTruth'][0,0]['Segmentation'][0][0]

    #target data also needs to be portrait
    if mat.shape != (481, 321):
        mat = numpy.rot90(mat)

    #pytorch expects single precision floats
    mat = mat.astype(numpy.float32)
    ground_truth.append(mat)
    
#combined input and target images
train_gt_set = []

#load ground truths and add to train set
for data, gt in zip(train_set,ground_truth):
    train_gt_set.append((data[0], gt))

#create a loader for automatic shuffling and batching
train_loader = torch.utils.data.DataLoader(train_gt_set, batch_size=40, shuffle=True)

#create our nn
net = Net()

#select loss function
criterion = nn.MSELoss()
#create optimiser
optimiser = optim.SGD(net.parameters(), lr=0.01)
#number of times data set will be fully passed through
num_epochs = 5

for epochs in range(num_epochs):
    for data in train_loader:
        inputs, segments = data

        #pytorch needs Variables to perform gradient calculations
        inputs, segments = Variable(inputs), Variable(segments)
        #reset gradients from previous iterations
        optimiser.zero_grad()
        #perform forward pass
        outputs = net(inputs)
        #calculate loss
        loss = criterion(outputs, segments)
        #perform back prop.
        loss.backward()
        #update weights
        optimiser.step()