initial commit

Author: gntoni

loadCOCO.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+import numpy as np
+import torch
+from os import path
+from scipy.misc import imread, imresize
+from scipy.io import loadmat
+
+
+class Rescale(object):
+    """Rescale the image in a sample to a given size.
+
+    Args:
+        output_size (int or tuple): Desired output size. If tuple, output is
+            matched to output_size. If int, smaller of image edges is matched
+            to output_size keeping aspect ratio the same.
+    """
+
+    def __init__(self, output_size):
+        assert isinstance(output_size, (int, tuple))
+        self.output_size = output_size
+
+    def __call__(self, image, labels):
+        assert image.shape[:2] == labels.shape
+
+        h, w = image.shape[:2]
+        if isinstance(self.output_size, int):
+            if h > w:
+                new_h, new_w = self.output_size * h / w, self.output_size
+            else:
+                new_h, new_w = self.output_size, self.output_size * w / h
+        else:
+            new_h, new_w = self.output_size
+
+        new_h, new_w = int(new_h), int(new_w)
+
+        img = imresize(image, (new_h, new_w))
+        lbls = imresize(labels, (new_h, new_w), interp="nearest")
+
+        return (img, lbls)
+
+
+class RandomCrop(object):
+    """Crop randomly the image in a sample.
+
+    Args:
+        output_size (tuple or int): Desired output size. If int, square crop
+            is made.
+    """
+
+    def __init__(self, output_size):
+        assert isinstance(output_size, (int, tuple))
+        if isinstance(output_size, int):
+            self.output_size = (output_size, output_size)
+        else:
+            assert len(output_size) == 2
+            self.output_size = output_size
+
+    def __call__(self, image, labels):
+        assert image.shape[:2] == labels.shape
+
+        h, w = image.shape[:2]
+        new_h, new_w = self.output_size
+
+        top = np.random.randint(0, h - new_h)
+        left = np.random.randint(0, w - new_w)
+
+        image = image[top: top + new_h,
+                      left: left + new_w]
+
+        labels = labels[top: top + new_h,
+                        left: left + new_w]
+
+        return (image, labels)
+
+
+class ToTensor(object):
+    """Convert ndarrays in sample to Tensors."""
+
+    def __call__(self, image, labels):
+        assert image.shape[:2] == labels.shape
+
+        # swap color axis because
+        # numpy image: H x W x C
+        # torch image: C X H X W
+        image = image.transpose((2, 0, 1))
+        return (torch.from_numpy(image),
+                torch.from_numpy(labels))
+
+
+def loadCOCO(dataset_folder):
+    resc = Rescale(650)
+    crop = RandomCrop(640)
+
+    namespath = path.join(dataset_folder, "imageLists/train.txt")
+    names = np.loadtxt(namespath, dtype=str, delimiter="\n")
+
+    images = []
+    labels = []
+    for imgName in names:
+        im = imread(path.join(dataset_folder, "images/"+imgName+".jpg"), mode="RGB")
+        mat = loadmat(path.join(dataset_folder, "annotations/"+imgName+".mat"))
+        lbl = mat["S"]
+
+        im, lbl = resc(im, lbl)
+        im, lbl = crop(im, lbl)
+        images.append(im)
+        labels.append(lbl)
+
+    images = np.array(images, dtype='float32')
+    images /= 255.0  # Span 0 ~ 1
+    images = (images*2) - 1  # Span -1 ~ 1
+
+    return (images, np.array(labels))
+
+
+if __name__ == '__main__':
+    DATASET_FOLDER = "/home/toni/Data/ssegmentation/COCO"
+    loadCOCO(DATASET_FOLDER)

unet.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv64 = nn.Conv2d(1, 64, 3, padding=1)
+        self.conv128 = nn.Conv2d(64, 128, 3, padding=1)
+        self.conv256 = nn.Conv2d(128, 256, 3, padding=1)
+        self.conv512 = nn.Conv2d(256, 512, 3, padding=1)
+        self.conv1024 = nn.Conv2d(512, 1024, 3, padding=1)
+        self.upconv1024 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
+        self.dconv1024 = nn.Conv2d(1024, 512, 3, padding=1)
+        self.upconv512 = nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.dconv512 = nn.Conv2d(512, 256, 3, padding=1)
+        self.upconv256 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.dconv256 = nn.Conv2d(256, 128, 3, padding=1)
+        self.upconv128 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.dconv128 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv1 = nn.Conv2d(64, 2, 1)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x1 = F.relu(self.conv64(x))
+        x2 = F.relu(self.conv128(self.pool(x1)))
+        x3 = F.relu(self.conv256(self.pool(x2)))
+        x4 = F.relu(self.conv512(self.pool(x3)))
+        x5 = F.relu(self.conv1024(self.pool(x4)))
+        ux5 = self.upconv1024(x5)
+        cc5 = torch.cat([ux5, x4], 1)
+        dx4 = F.relu(self.dconv1024(cc5))
+        ux4 = self.upconv512(dx4)
+        cc4 = torch.cat([ux4, x3], 1)
+        dx3 = F.relu(self.dconv512(cc4))
+        ux3 = self.upconv256(dx3)
+        cc3 = torch.cat([ux3, x2], 1)
+        dx2 = F.relu(self.dconv256(cc3))
+        ux2 = self.upconv128(dx2)
+        cc2 = torch.cat([ux2, x1], 1)
+        dx1 = F.relu(self.dconv128(cc2))  # no relu?
+        last = self.conv1(dx1)
+        return F.log_softmax(last)  # sigmoid if classes arent mutually exclusv
+
+###########
+# Load Dataset  #
+###########
+
+
+net = Net()
+criterion = nn.NLLLoss2d()
+optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+
+for epoch in range(2):  # loop over the dataset multiple times
+    running_loss = 0.0
+    for i, data in enumerate(trainloader, 0):
+        # get the inputs
+        inputs, labels = data
+
+        # wrap them in Variable
+        inputs, labels = Variable(inputs), Variable(labels)
+
+        # zero the parameter gradients
+        optimizer.zero_grad()
+
+        # forward + backward + optimize
+        outputs = net(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        # print statistics
+        running_loss += loss.data[0]
+        if i % 2000 == 1999:    # print every 2000 mini-batches
+            print('[%d, %5d] loss: %.3f' %
+                  (epoch + 1, i + 1, running_loss / 2000))
+            running_loss = 0.0
+
+print('Finished Training')