ml.py

# -*- coding: utf-8 -*-
"""ML.ipynb

Automatically generated by Colaboratory.

Make sure your directory structure is as follows:
- ./src/img contains all the training images
- create directory ./src/lbl which is initially empty or holds training points
- If you want to add more training data then uncomment lines 45-47

Original file is located at
    https://colab.research.google.com/drive/1zGkpNkV_v0B1f91L2g_TU6eu1uJ_R5Ys
"""

import os

import tqdm
import skimage.io as skio
import skimage.transform as sktr
import matplotlib.pyplot as plt
import numpy as np

os.listdir("./src/img/")

def label_image(im_name, n=8):
    im = skio.imread("./src/img/" + im_name)
    im2 = sktr.resize(im, (300, 500, 3))
    skio.imsave("./src/img/" + im_name, im2)

    try:
        np.load("./src/lbl/" + im_name[:-4] + ".npy")
    except:
        print("Please select " + str(n) + "points")

        plt.imshow(im)
        points = plt.ginput(n, mouse_stop=None, timeout=-1)
        plt.close()

        points = np.array(points) / (im.shape[1], im.shape[0])

        np.save("./src/lbl/" + im_name[:-4], points)

# Make sure that your directory structure is as follows -
# %matplotlib qt
# for file in os.listdir("./src/img/"):
#     print(file)
#     label_image(file)
# %matplotlib inline

import torch
from torch.utils.data import Dataset, Subset
import skimage.io as io
from torch import nn
import cv2
import numpy as np
import os 
from skimage import exposure


def magnitude(dx, dy):
    return np.sqrt(dx ** 2 + dy **2)

class Table(Dataset):
    """Table Landmarks dataset."""

    def __init__(self, root_dir, transform=None):
        """
        Args:
            root_dir (string): Directory with all the data.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.root_dir = root_dir
        self.images = os.listdir(root_dir + "/img/")
        self.points = os.listdir(root_dir + "/lbl/")
        self.len = len(self.images)
        # self.c = ColorJitter(brightness=.15, contrast=.15)
        
        self.labels = {root_dir + "/img/" + b[:-4] + ".png" : root_dir + "/lbl/" + b for b in self.points}

        self.transformation = transform

    def __len__(self):
        return self.len * 12
  
    def transform(self, sample):
        sample['image'] = exposure.equalize_adapthist(sample['image'], clip_limit=0.005)
        # pil_image = Image.fromarray((sample['image'] * 255).astype('uint8'), 'L')
        # jittered_pil_image = self.c(pil_image)
        # jittered_np_image = (np.array(jittered_pil_image).astype(np.float32) / 255 - 0.5)
        # sample['image'] = jittered_np_image
        # k = cv2.getGaussianKernel(6, 5)

        # kernel = k @ k.T
        # kernel = convolve2d(kernel, kernel)
        # kernel = convolve2d(kernel, kernel)
        # d_x = np.array([1, -1]).reshape(1, 2)
        # d_y = np.array([1, -1]).reshape(2, 1)
        # derivative_x_kernel = convolve2d(kernel, d_x)
        # derivative_y_kernel = convolve2d(kernel, d_y)


        p = np.random.random()
        if p < 0.7:
            sample = self.rotate(sample)
        p = np.random.random()
        if p < 0.7:
            sample = self.translate(sample)

        # im = sample['image']

        # blurred_x_derivative = convolve2d(im, derivative_x_kernel, mode="same", boundary="wrap")
        # blurred_y_derivative = convolve2d(im, derivative_y_kernel, mode="same", boundary="wrap")
        # s, t = np.percentile(magnitude(blurred_x_derivative, blurred_y_derivative), [97,100])
        # sample['image'] = magnitude(blurred_x_derivative, blurred_y_derivative) > s

        return sample

    def rotate(self, sample):
        image = sample['image']
        y, x = tuple(np.array(image.shape[1::-1]) / 2)
        y, x = np.random.random(), np.random.random()
        y, x = y * np.random.randint(y-2, y+2), np.random.randint(x-2, x+2)
        image_center = (y, x)
        angle = (np.random.random()) * 4
        rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
        result = cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR)
        sample['image'] = result
        sample["landmarks"] = self.rotate_points(sample, sample["landmarks"], rot_mat)
        return sample

    def rotate_points(self, sample, pts, rot_mat):
        image = sample['image']
        # transform points
        points = pts * image.shape[1::-1]
        # add ones
        ones = np.ones(shape=(len(points), 1))
        points_ones = np.hstack([points, ones])
        # transform points
        transformed_points = rot_mat.dot(points_ones.T).T / image.shape[1::-1]
        return transformed_points

    def translate(self, sample):
        image = sample['image']

        # Store height and width of the image
        height, width = image.shape[:2]
          
        y, x = np.random.randint(4, size=2)
          
        T = np.float32([[1, 0, x], [0, 1, y]])
          
        # We use warpAffine to transform
        # the image using the matrix, T
        img_translation = cv2.warpAffine(image, T, (width, height))

        sample['image'] = img_translation

        points = sample["landmarks"] * image.shape[1::-1]
        points[:, 0] += x
        points[:, 1] += y
        points = points / image.shape[1::-1]
        sample["landmarks"] = points

        return sample

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
            
        idx = idx % self.len

        img_name = self.root_dir + "/img/" + self.images[idx]
        sample = {'image': io.imread(img_name)[..., 1], 'landmarks': np.load(self.labels[img_name])}

        sample = self.transform(sample)

        return sample

dataset = Table(root_dir='./src')

for a in dataset.images:
    print(a[:-4])

def show_landmarks(ax, image, landmarks):
    """Show image with landmarks"""

    height, width = image.shape
    ax.imshow(image, cmap='gray')
    landmarks = landmarks * (width, height)
    x, y = landmarks.T
    ax.scatter(x, y, s=10, marker='o', c='r')


length = len(os.listdir("./src/lbl"))
idxs = np.arange(length)
np.random.shuffle(idxs)

train_idxs, val_idxs = idxs[:int(length * 0.8)], idxs[int(length * 0.8):]

expanded_train = train_idxs.copy()
for i in train_idxs:
  for j in range(1, 12):
    expanded_train = np.append(expanded_train, i + j * 10)

expanded_val = val_idxs.copy()
for j in range(1, 12):
  for i in val_idxs:
    expanded_val = np.append(expanded_val, i + j * 10)

# train_idxs, val_idxs
train_data = Subset(dataset, expanded_train)
val_data = Subset(dataset, expanded_val)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 24, 3, 1, padding=2),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(24, 24, 3, 1, padding=2),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(24, 32, 3, 1, padding=2),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(32, 32, 3, 1),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(32, 32, 2, 1, padding=2),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(32, 64, 2, 1),
            nn.ReLU(),
        )

        self.decoder = nn.Sequential(
            nn.Linear(9216, 5120),
            nn.ReLU(),
            nn.Linear(5120, 16),
            nn.Sigmoid()
        )

    def forward(self, x):
        batch_size = x.shape[0]
        features = self.encoder(x).squeeze()

        product = np.prod(np.array(features.shape))
        features = features.reshape(batch_size, product // batch_size)

        return self.decoder(features).reshape(batch_size, 8, 2)

def accuracy(x, y):
    num_correct = 0
    for pred, real in zip(x, y):
        diffs = pred.cpu().detach().numpy() - real.cpu().detach().numpy()
        norms = np.linalg.norm(diffs, axis=1)
        num_correct += sum(norms < .05) / len(norms)
    return num_correct

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device", device)

epochs = 20
batch_size = 12
learning_rate = 0.0003
model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()
dataloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)

# model.train() # Put model in training mode
training_accuracies = []
validation_accuracies = []
t_loss = []
v_loss = []
for epoch in range(epochs):
    model.train()
    training_losses = []
    val_losses = []
    num_correct = 0
    for datapoint in tqdm.tqdm(dataloader, unit="batch"):
        x, y = datapoint['image'], datapoint['landmarks']
        x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
        x, y = x.float().to(device), y.float().to(device)
        # optimizer.zero_grad() # Remove the gradients from the previous step
        pred = model(x)
        num_correct += accuracy(pred, y)
        loss = criterion(pred, y)
        optimizer.zero_grad() # Remove the gradients from the previous step
        loss.backward()
        optimizer.step()
        training_losses.append(loss.item())
    training_accuracies.append(num_correct/len(train_data))
    t_loss.append(np.mean(training_losses))
    model.train() # Put model back in train mode
    print("Finished Epoch", epoch + 1, ", training loss:", np.mean(training_losses), ", training accuracy:", training_accuracies[-1])
    
    with torch.no_grad():
        val_dataloader = torch.utils.data.DataLoader(val_data)
        model.eval() # Put model in eval mode
        num_correct = 0
        preds = []
        validation_losses = []
        for datapoint in val_dataloader:
            x, y = datapoint['image'], datapoint['landmarks']
            x, y = x.float().to(device), y.float().to(device)
            x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
            pred = model(x)
            num_correct += accuracy(pred, y)
            loss = criterion(pred, y)
            validation_losses.append(loss.item())
        v_loss.append(np.mean(validation_losses))
        validation_accuracies.append(num_correct/len(val_data))
        print("Validation_Accuracy", validation_accuracies[-1])

with torch.no_grad():
        val_dataloader = torch.utils.data.DataLoader(train_data)
        model.eval() # Put model in eval mode
        num_correct = 0
        preds = []
        for datapoint in val_dataloader:
            x, y = datapoint['image'], datapoint['landmarks']
            x, y = x.float().to(device), y.float().to(device)
            x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
            pred = model(x)
            # preds.append(torch.argmax(pred, 0).item())
            img = datapoint['image']
            a, b = pred.cpu().T
            c, d = y.cpu().T
            a = a * img.shape[2]
            b = b * img.shape[1]
            c = c * img.shape[2]
            d = d * img.shape[1]
            num_correct += accuracy(pred, y)
        # preds = np.array(preds)
        # print(np.unique(preds, return_counts=True))
        validation_accuracies.append(num_correct/len(val_data))
        print("Validation_Accuracy", validation_accuracies[-1])

torch.save(model.state_dict(), "./dict.pth")
# torch.save(model.cpu(), "./model_2.pth")
#
# import torchvision.models as models
# def large_model():
#     model = models.resnet34(pretrained=True)
#     model.conv1 = nn.Conv2d(1, 64, 7, stride=2, padding=3, bias=False)
#     model.fc = nn.Linear(512, 16)
#     return model
#
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print("Using device", device)
#
# epochs = 20
# batch_size = 24
# learning_rate = 0.0005
# transfer_model = large_model().to(device)
# optimizer = torch.optim.Adam(transfer_model.parameters(), lr=learning_rate)
# criterion = nn.MSELoss()
# dataloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True)
#
# # model.train() # Put model in training mode
# training_accuracies = []
# validation_accuracies = []
# t_loss = []
# v_loss = []
# for epoch in range(epochs):
#     transfer_model.train()
#     training_losses = []
#     val_losses = []
#     num_correct = 0
#     for datapoint in tqdm.tqdm(dataloader, unit="batch"):
#         x, y = datapoint['image'], datapoint['landmarks']
#         x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
#         x, y = x.float().to(device), y.float().to(device)
#         # optimizer.zero_grad() # Remove the gradients from the previous step
#         pred = transfer_model(x).view((x.shape[0], 8, 2))
#         num_correct += accuracy(pred, y)
#         loss = criterion(pred, y)
#         optimizer.zero_grad() # Remove the gradients from the previous step
#         loss.backward()
#         optimizer.step()
#         training_losses.append(loss.item())
#     training_accuracies.append(num_correct/len(train_data))
#     t_loss.append(np.mean(training_losses))
#     transfer_model.train() # Put model back in train mode
#     print("Finished Epoch", epoch + 1, ", training loss:", np.mean(training_losses), ", training accuracy:", training_accuracies[-1])
#
#     with torch.no_grad():
#         val_dataloader = torch.utils.data.DataLoader(val_data)
#         transfer_model.eval() # Put model in eval mode
#         num_correct = 0
#         preds = []
#         validation_losses = []
#         for datapoint in val_dataloader:
#             x, y = datapoint['image'], datapoint['landmarks']
#             x, y = x.float().to(device), y.float().to(device)
#             x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
#             pred = transfer_model(x).view((1, 8, 2))
#             num_correct += accuracy(pred, y)
#             loss = criterion(pred, y)
#             validation_losses.append(loss.item())
#         v_loss.append(np.mean(validation_losses))
#         validation_accuracies.append(num_correct/len(val_data))
#         print("Validation_Accuracy", validation_accuracies[-1])
#
# transfer_model =torch.load("./transfer.pth")
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#
# with torch.no_grad():
#         val_dataloader = torch.utils.data.DataLoader(val_data)
#         transfer_model.eval() # Put model in eval mode
#         num_correct = 0
#         preds = []
#         for datapoint in val_dataloader:
#             x, y = datapoint['image'], datapoint['landmarks']
#             x, y = x.float().to(device), y.float().to(device)
#             x = torch.reshape(x, (x.shape[0], 1, x.shape[1], x.shape[2]))
#             pred = transfer_model(x)
#             # preds.append(torch.argmax(pred, 0).item())
#             img = datapoint['image']
#             a, b = pred.view((1, 8, 2)).cpu().T
#             c, d = y.cpu().T
#             a = a * img.shape[2]
#             b = b * img.shape[1]
#             c = c * img.shape[2]
#             d = d * img.shape[1]
#             num_correct += accuracy(pred.view(1, 8, 2), y)
#         # preds = np.array(preds)
#         # print(np.unique(preds, return_counts=True))
#         validation_accuracies.append(num_correct/len(val_data))
#         print("Validation_Accuracy", validation_accuracies[-1])
#
# torch.save(transfer_model, "./transfer.pth")