Udacity - AI programming with python

predict.py

#!/usr/bin/env python3

# Example: python predict.py ../aipnd-project/flowers/test/10/image_07090.jpg ../model/checkpoint.pth
# Example: python predict.py ../aipnd-project/flowers/test/10/image_07090.jpg ../model/checkpoint.pth --topk 5 --gpu
# Example: python predict.py ../aipnd-project/flowers/test/10/image_07090.jpg ../model/checkpoint.pth --topk 5 --category_names ../aipnd-project/cat_to_name.json --gpu

# https://github.com/pytorch/examples/blob/master/imagenet/main.py#L139

import argparse

import json
import numpy as np
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision.models as models
from torch import nn
from PIL import Image

def get_args():
    args = {}
   
    parser = argparse.ArgumentParser(description='Image prediction')
    parser.add_argument("input")
    parser.add_argument("checkpoint")
    parser.add_argument('--topk', type=int, default=5)
    parser.add_argument('--category_names', type=str, default=None)
    parser.add_argument('--gpu', action='store_true')
    parsed_args = parser.parse_args()
    
    args = vars(parsed_args)

    return args

def load_model(checkpoint_path, gpu):
    model = models.vgg16(pretrained=False)
    
    classifier = nn.Sequential(OrderedDict([ 
        ('fc1', nn.Linear(25088, 2000)), 
        ('relu', nn.ReLU()), 
        ('fc2', nn.Linear(2000, 500)), 
        ('relu', nn.ReLU()), 
        ('fc3', nn.Linear(500, 102)), 
        ('output', nn.LogSoftmax(dim=1))]))
    
    model.classifier = classifier
    
    if(gpu == True):
        model.cuda()
    
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint["state_dict"])
    model.class_to_idx = checkpoint["class_to_idx"]
    
    return model

def process_image(image):  
    size = 256, 256
    
    new_side = 224
    center = 256 / 2
    start = 0 #center - (new_side / 2)
    
    im = Image.open(image)
    im.thumbnail(size)
    im = im.crop((start, start, new_side, new_side))
    np_image = np.array(im) / 255
    #print(np_image.shape)
    
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    np_image = (np_image - mean) / std
    
    #np_image = np_image.transpose((1, 2, 0))
    np_image = np.transpose(np_image, (2, 0, 1))
    
    return np_image

def predict(image_path, model, topk):
    im = process_image(image_path)
    input = torch.FloatTensor(im).cuda()
    input.unsqueeze_(0)
    output = model.forward(input)
    ps = F.softmax(output, dim=1)
                          
    return torch.topk(ps, topk)

def get_cat_to_name(cat_to_name_path):
    cat_to_name = []
    
    with open(cat_to_name_path, 'r') as f:
        cat_to_name = json.load(f)

    return cat_to_name
        
def main():
    args = get_args()
    model = load_model(args["checkpoint"], args["gpu"])

    probs, classes = predict(args["input"], model, args["topk"])

    print("Most likely classes: \n")
    
    if(args["category_names"] == None):
        print(classes[0])
    else:    
        cat_to_name = get_cat_to_name(args["category_names"])
        
        for c in classes[0]:  
            for m in model.class_to_idx:
                if model.class_to_idx[m] == c:
                    print("* " + cat_to_name[m])
                
if __name__ == "__main__":
    main()

train.py

#!/usr/bin/env python3

# Example: python train.py ../aipnd-project/flowers --arch vgg --gpu
# Example: python train.py ../aipnd-project/flowers --arch vgg --save_dir ../model --learning_rate 0.00005 --hidden_units 500 --epochs 3 --gpu

import sys
import argparse

from collections import OrderedDict

import torch
import torch.nn.functional as F
from torch.autograd import Variable
from torch import optim
from torch import nn
import torchvision.models as models
from torchvision import datasets, transforms

resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

models = {'resnet': resnet18, 'alexnet': alexnet, 'vgg': vgg16}

def get_args():
    args = {}
   
    parser = argparse.ArgumentParser(description='Image classifier')
    parser.add_argument("image_dir")
    parser.add_argument('--arch', type=str, choices=["vgg", "alexnet", "resnet"], required=True)
    parser.add_argument('--learning_rate', type=float, default=0.00005)
    parser.add_argument('--hidden_units', type=int, default=500)
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--gpu', action='store_true')
    parser.add_argument('--save_dir', type=str, default="../model")
    parsed_args = parser.parse_args()
    
    args = vars(parsed_args)

    return args

def get_dataloaders(train_dir, test_dir):
    data_transforms = transforms.Compose([transforms.RandomRotation(30),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], 
        [0.229, 0.224, 0.225])])

    dataloaders = {}
    
    trainset = datasets.ImageFolder(train_dir, transform=data_transforms)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
    
    testset = datasets.ImageFolder(test_dir, transform=data_transforms)
    testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=True)
    
    dataloaders["train"] = {
        "dataset": trainset,
        "loader": trainloader
    }    
    
    dataloaders["test"] = {
        "dataset": testset,
        "loader": testloader
    }    
    
    return dataloaders

def create_model(arch):
    model = models[arch]
    
    return model

def train_model(model, trainloader, testloader, learning_rate, hidden_units, epochs, gpu):
    classifier = nn.Sequential(OrderedDict([ 
        ('fc1', nn.Linear(25088, 2000)), 
        ('relu', nn.ReLU()), 
        ('fc2', nn.Linear(2000, hidden_units)), 
        ('relu', nn.ReLU()), 
        ('fc3', nn.Linear(hidden_units, 102)), 
        ('output', nn.LogSoftmax(dim=1))]))
    
    model.classifier = classifier
    
    if(gpu == True):
        model.cuda()

    steps = 0
    running_loss = 0
    print_every = 40

    criterion = nn.NLLLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for e in range(epochs):
        # Model in training mode, dropout is on
        model.train()
    
        for images, labels in iter(trainloader):
            steps += 1
            # Flatten images into a 784 long vector
            #images.resize_(images.size()[0], 224)
        
            # Wrap images and labels in Variables so we can calculate gradients
            inputs = Variable(images)
            targets = Variable(labels)
            optimizer.zero_grad()
        
            if(gpu == True): 
                inputs, labels = inputs.cuda(), labels.cuda()
                targets = targets.cuda()
        
            output = model.forward(inputs)
        
            if(gpu == True):
                output = output.cuda()
                
            loss = criterion(output, targets)
            loss.backward()
            optimizer.step()
        
            running_loss += loss.data[0]
        
            if steps % print_every == 0:
                # Model in inference mode, dropout is off
                model.eval()
                
                if(gpu == True):
                    model.cuda()
            
                accuracy = 0
                test_loss = 0
                
                for ii, (images, labels) in enumerate(testloader):
                    #images = images.resize_(images.size()[0], 500)
                    # Set volatile to True so we don't save the history
                    inputs = Variable(images, volatile=True)
                    labels = Variable(labels, volatile=True)
                
                    if(gpu == True):
                        inputs, labels = inputs.cuda(), labels.cuda()

                    output = model.forward(inputs)
                    test_loss += criterion(output, labels).data[0]
                
                    ## Calculating the accuracy 
                    # Model's output is log-softmax, take exponential to get the probabilities
                    ps = torch.exp(output).data
                    # Class with highest probability is our predicted class, compare with true label
                    equality = (labels.data == ps.max(1)[1])
                    # Accuracy is number of correct predictions divided by all predictions, just take the mean
                    accuracy += equality.type_as(torch.FloatTensor()).mean()
            
                    print("Epoch: {}/{}.. ".format(e+1, epochs),
                        "Training Loss: {:.3f}.. ".format(running_loss/print_every),
                        "Test Loss: {:.3f}.. ".format(test_loss/len(testloader)),
                        "Test Accuracy: {:.3f}".format(accuracy/len(testloader)))
            
                running_loss = 0
            
                # Make sure dropout is on for training
                model.train()
                
    return model

def save_model(model, trainset, save_dir):
    model.class_to_idx = trainset.class_to_idx
    torch.save({
        'arch': 'vgg16',
        'state_dict': model.state_dict(), 
        'class_to_idx': model.class_to_idx
    }, save_dir + '/checkpoint.pth')

def main():
    args = get_args()
    
    train_dir = args["image_dir"] + "/train"
    test_dir = args["image_dir"] + "/test"
    dataloaders = get_dataloaders(train_dir, test_dir)
    
    model = create_model(args["arch"])
    model = train_model(
        model, 
        dataloaders["train"]["loader"], 
        dataloaders["test"]["loader"],
        args["learning_rate"],
        args["hidden_units"],
        args["epochs"],
        args["gpu"])
    
    save_model(model, dataloaders["train"]["dataset"], args["save_dir"])
    
if __name__ == "__main__":
    main()