<reponame>manusimidt/udacity_p1_facial_keypoints
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
from models import Net
from data_load import FacialKeypointsDataset, Rescale, RandomCrop, Normalize, ToTensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")
print(f"Using device {device}: {device}")
neural_net = Net().to(device)
print(neural_net)

data_transform = transforms.Compose([Rescale(250), RandomCrop(224), Normalize(), ToTensor()])
# create the transformed dataset
transformed_dataset = FacialKeypointsDataset(csv_file='data/training_frames_keypoints.csv',
                                             root_dir='data/training/',
                                             transform=data_transform)
test_dataset = FacialKeypointsDataset(csv_file='data/test_frames_keypoints.csv', root_dir='data/test/',
                                      transform=data_transform)
train_loader = DataLoader(transformed_dataset, batch_size=25, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=25, shuffle=True, num_workers=4)

criterion = nn.MSELoss()
optimizer = optim.Adam(params=neural_net.parameters(), lr=0.001)


def net_sample_output():
    # iterate through the test dataset
    for i, sample in enumerate(test_loader):

        # get sample data: images and ground truth keypoints
        images = sample['image']
        key_pts = sample['keypoints']

        # convert images to FloatTensors
        # images = images.type(torch.cuda.FloatTensor)
        images = images.float().to(device)

        # forward pass to get net output
        output_pts = neural_net(images)

        # reshape to batch_size x 68 x 2 pts
        output_pts = output_pts.view(output_pts.size()[0], 68, -1)

        # break after first image is tested
        if i == 0:
            return images, output_pts, key_pts


def show_all_keypoints(image, predicted_key_pts, gt_pts=None):
    """Show image with predicted keypoints"""
    # image is grayscale
    plt.imshow(image, cmap='gray')
    plt.scatter(predicted_key_pts[:, 0], predicted_key_pts[:, 1], s=20, marker='.', c='m')
    # plot ground truth points as green pts
    if gt_pts is not None:
        plt.scatter(gt_pts[:, 0], gt_pts[:, 1], s=20, marker='.', c='g')


def visualize_output(test_images, test_outputs, gt_pts=None, batch_size=10):
    for i in range(batch_size):
        plt.figure(figsize=(20, 10))
        ax = plt.subplot(1, batch_size, i + 1)

        # un-transform the image data
        image = test_images[i].data  # get the image from it's Variable wrapper
        image = image.cpu().numpy()  # convert to numpy array from a Tensor
        image = np.transpose(image, (1, 2, 0))  # transpose to go from torch to numpy image

        # un-transform the predicted key_pts data
        predicted_key_pts = test_outputs[i].data
        predicted_key_pts = predicted_key_pts.cpu().numpy()
        # undo normalization of keypoints
        predicted_key_pts = predicted_key_pts * 50.0 + 100

        # plot ground truth points for comparison, if they exist
        ground_truth_pts = None
        if gt_pts is not None:
            ground_truth_pts = gt_pts[i]
            ground_truth_pts = ground_truth_pts * 50.0 + 100

        # call show_all_keypoints
        show_all_keypoints(np.squeeze(image), predicted_key_pts, ground_truth_pts)

        plt.axis('off')

    plt.show()


def train_net(n_epochs: int) -> None:
    """
    Trains the network
    """
    neural_net.train()

    for epoch in range(n_epochs):  # loop over the dataset multiple times

        running_loss = 0.0

        # train on batches of data, assumes you already have train_loader
        for batch_i, data in enumerate(train_loader):
            # get the input images and their corresponding labels
            images = data['image']
            key_pts = data['keypoints']

            # flatten pts
            key_pts = key_pts.view(key_pts.size(0), -1)

            # convert variables to floats for regression loss
            # key_pts = key_pts.type(torch.cuda.FloatTensor)
            # images = images.type(torch.cuda.FloatTensor)
            key_pts = key_pts.float().to(device)
            images = images.float().to(device)
            # forward pass to get outputs
            output_pts = neural_net.forward(images)
            # calculate the loss between predicted and target keypoints
            loss = criterion(output_pts, key_pts)
            # zero the parameter (weight) gradients
            optimizer.zero_grad()
            # backward pass to calculate the weight gradients
            loss.backward()
            # update the weights
            optimizer.step()
            # print loss statistics
            running_loss += loss.item()

            torch.cuda.empty_cache()

            if batch_i % 10 == 9:  # print every 10 batches
                print('Epoch: {}, Batch: {}, Avg. Loss: {}'.format(epoch + 1, batch_i + 1, running_loss / 10))
                running_loss = 0.0
    print('Finished Training')


if __name__ == '__main__':
    train_net(10)
    visualize_output(*net_sample_output())

    # after training, save your model parameters in the dir 'saved_models'
    torch.save(neural_net.state_dict(), 'saved_models/' + 'model1')

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
from models import Net
from data_load import FacialKeypointsDataset, Rescale, RandomCrop, Normalize, ToTensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")
print(f"Using device {device}: {device}")
neural_net = Net().to(device)
print(neural_net)

data_transform = transforms.Compose([Rescale(250), RandomCrop(224), Normalize(), ToTensor()])
# create the transformed dataset
transformed_dataset = FacialKeypointsDataset(csv_file='data/training_frames_keypoints.csv',
                                             root_dir='data/training/',
                                             transform=data_transform)
test_dataset = FacialKeypointsDataset(csv_file='data/test_frames_keypoints.csv', root_dir='data/test/',
                                      transform=data_transform)
train_loader = DataLoader(transformed_dataset, batch_size=25, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=25, shuffle=True, num_workers=4)

criterion = nn.MSELoss()
optimizer = optim.Adam(params=neural_net.parameters(), lr=0.001)


def net_sample_output():
    # iterate through the test dataset
    for i, sample in enumerate(test_loader):

        # get sample data: images and ground truth keypoints
        images = sample['image']
        key_pts = sample['keypoints']

        # convert images to FloatTensors
        # images = images.type(torch.cuda.FloatTensor)
        images = images.float().to(device)

        # forward pass to get net output
        output_pts = neural_net(images)

        # reshape to batch_size x 68 x 2 pts
        output_pts = output_pts.view(output_pts.size()[0], 68, -1)

        # break after first image is tested
        if i == 0:
            return images, output_pts, key_pts


def show_all_keypoints(image, predicted_key_pts, gt_pts=None):
    """Show image with predicted keypoints"""
    # image is grayscale
    plt.imshow(image, cmap='gray')
    plt.scatter(predicted_key_pts[:, 0], predicted_key_pts[:, 1], s=20, marker='.', c='m')
    # plot ground truth points as green pts
    if gt_pts is not None:
        plt.scatter(gt_pts[:, 0], gt_pts[:, 1], s=20, marker='.', c='g')


def visualize_output(test_images, test_outputs, gt_pts=None, batch_size=10):
    for i in range(batch_size):
        plt.figure(figsize=(20, 10))
        ax = plt.subplot(1, batch_size, i + 1)

        # un-transform the image data
        image = test_images[i].data  # get the image from it's Variable wrapper
        image = image.cpu().numpy()  # convert to numpy array from a Tensor
        image = np.transpose(image, (1, 2, 0))  # transpose to go from torch to numpy image

        # un-transform the predicted key_pts data
        predicted_key_pts = test_outputs[i].data
        predicted_key_pts = predicted_key_pts.cpu().numpy()
        # undo normalization of keypoints
        predicted_key_pts = predicted_key_pts * 50.0 + 100

        # plot ground truth points for comparison, if they exist
        ground_truth_pts = None
        if gt_pts is not None:
            ground_truth_pts = gt_pts[i]
            ground_truth_pts = ground_truth_pts * 50.0 + 100

        # call show_all_keypoints
        show_all_keypoints(np.squeeze(image), predicted_key_pts, ground_truth_pts)

        plt.axis('off')

    plt.show()


def train_net(n_epochs: int) -> None:
    """
    Trains the network
    """
    neural_net.train()

    for epoch in range(n_epochs):  # loop over the dataset multiple times

        running_loss = 0.0

        # train on batches of data, assumes you already have train_loader
        for batch_i, data in enumerate(train_loader):
            # get the input images and their corresponding labels
            images = data['image']
            key_pts = data['keypoints']

            # flatten pts
            key_pts = key_pts.view(key_pts.size(0), -1)

            # convert variables to floats for regression loss
            # key_pts = key_pts.type(torch.cuda.FloatTensor)
            # images = images.type(torch.cuda.FloatTensor)
            key_pts = key_pts.float().to(device)
            images = images.float().to(device)
            # forward pass to get outputs
            output_pts = neural_net.forward(images)
            # calculate the loss between predicted and target keypoints
            loss = criterion(output_pts, key_pts)
            # zero the parameter (weight) gradients
            optimizer.zero_grad()
            # backward pass to calculate the weight gradients
            loss.backward()
            # update the weights
            optimizer.step()
            # print loss statistics
            running_loss += loss.item()

            torch.cuda.empty_cache()

            if batch_i % 10 == 9:  # print every 10 batches
                print('Epoch: {}, Batch: {}, Avg. Loss: {}'.format(epoch + 1, batch_i + 1, running_loss / 10))
                running_loss = 0.0
    print('Finished Training')


if __name__ == '__main__':
    train_net(10)
    visualize_output(*net_sample_output())

    # after training, save your model parameters in the dir 'saved_models'
    torch.save(neural_net.state_dict(), 'saved_models/' + 'model1')
# device = torch.device("cpu") # create the transformed dataset # iterate through the test dataset # get sample data: images and ground truth keypoints # convert images to FloatTensors # images = images.type(torch.cuda.FloatTensor) # forward pass to get net output # reshape to batch_size x 68 x 2 pts # break after first image is tested Show image with predicted keypoints # image is grayscale # plot ground truth points as green pts # un-transform the image data # get the image from it's Variable wrapper # convert to numpy array from a Tensor # transpose to go from torch to numpy image # un-transform the predicted key_pts data # undo normalization of keypoints # plot ground truth points for comparison, if they exist # call show_all_keypoints Trains the network # loop over the dataset multiple times # train on batches of data, assumes you already have train_loader # get the input images and their corresponding labels # flatten pts # convert variables to floats for regression loss # key_pts = key_pts.type(torch.cuda.FloatTensor) # images = images.type(torch.cuda.FloatTensor) # forward pass to get outputs # calculate the loss between predicted and target keypoints # zero the parameter (weight) gradients # backward pass to calculate the weight gradients # update the weights # print loss statistics # print every 10 batches # after training, save your model parameters in the dir 'saved_models'