Completed notebook, exported functions to respective files

.gitignore

@@ -99,3 +99,6 @@ ENV/
 
 # mypy
 .mypy_cache/
+
+data/
+*.pth

helper.py

@@ -0,0 +1,80 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+
+def imshow_original(image, ax=None, title=None, normalize=True):
+    """Imshow for Tensor."""
+    if ax is None:
+        fig, ax = plt.subplots()
+    image = image.numpy().transpose((1, 2, 0))
+
+    if normalize:
+        mean = np.array([0.485, 0.456, 0.406])
+        std = np.array([0.229, 0.224, 0.225])
+        image = std * image + mean
+        image = np.clip(image, 0, 1)
+
+    ax.imshow(image)
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.spines['left'].set_visible(False)
+    ax.spines['bottom'].set_visible(False)
+    ax.tick_params(axis='both', length=0)
+    ax.set_xticklabels('')
+    ax.set_yticklabels('')
+
+    return ax
+
+def process_image(img_path):
+    ''' Scales, crops, and normalizes a PIL image for a PyTorch model,
+        returns an Numpy array
+    '''
+    img = Image.open(img_path)
+    w, h = img.size
+    if w<h:
+        size = 256, 999999999
+    else:
+        size = 999999999, 256
+
+    img.thumbnail(size=size)
+
+    w, h = img.size
+    left = (w - 224) / 2
+    right = (w + 224) / 2
+    top = (h - 224) / 2
+    bottom = (h + 224) / 2
+
+    img = img.crop((left, top, right, bottom))
+
+    # Convert to numpy array
+    np_img = np.array(img)/255
+
+    # Normalize
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    np_img = (np_img - mean) / std
+
+    np_img = np_img.transpose(2, 0, 1)
+
+    return np_img
+
+def imshow(image, ax=None, title=None):
+    """Imshow for Tensor."""
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    # PyTorch tensors assume the color channel is the first dimension
+    # but matplotlib assumes is the third dimension
+    image = image.transpose((1, 2, 0))
+
+    # Undo preprocessing
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    image = std * image + mean
+
+    # Image needs to be clipped between 0 and 1 or it looks like noise when displayed
+    image = np.clip(image, 0, 1)
+
+    ax.imshow(image)
+
+    return ax

predict.py

@@ -0,0 +1,52 @@
+import torch
+import numpy as np
+from torchvision import models
+from helper import process_image
+
+def load_model(model_checkpoint):
+    checkpoint = torch.load(model_checkpoint)
+
+    model = models.vgg16(pretrained=True)
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    model.classifier = checkpoint["classifier"]
+    model.load_state_dict(checkpoint["state_dict"])
+    class_idx_mapping = checkpoint["class_idx_mapping"]
+    idx_class_mapping = {v: k for k, v in class_idx_mapping.items()}
+
+    return model, idx_class_mapping
+
+def predict(image_path, model_checkpoint, topk=5):
+    ''' 
+    Predict the class (or classes) of an image using a trained deep learning model.
+    
+    Arguments:
+        image_path: Path to the image
+        model: Trained model
+    '''
+
+    # Build the model from the checkpoint
+    model, idx_class_mapping = load_model(model_checkpoint)
+
+    # No need for GPU
+    model.to("cpu")
+
+    model.eval()
+
+    img = process_image(image_path)
+    img = np.expand_dims(img, axis=0)
+    img_tensor = torch.from_numpy(img).type(torch.FloatTensor).to('cpu')
+
+    with torch.no_grad():
+        log_probabilities = model.forward(img_tensor)
+
+    probabilities = torch.exp(log_probabilities)
+    probs, indices = probabilities.topk(5)
+
+    probs = probs.numpy().squeeze()
+    indices = indices.numpy().squeeze()
+    classes = [idx_class_mapping[index] for index in indices]
+
+    return probs, classes

train.py

@@ -0,0 +1,82 @@
+import torch
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+def validation(model, testloader, criterion):
+    test_loss = 0
+    accuracy = 0
+    model.to(device)
+    for images, labels in testloader:
+        images, labels = images.to(device), labels.to(device)
+        # images.resize_(images.shape[0], 3, 224, 224)
+
+        output = model.forward(images)
+        test_loss += criterion(output, labels).item()
+
+        ps = torch.exp(output)
+        equality = (labels.data == ps.max(dim=1)[1])
+        accuracy += equality.type(torch.FloatTensor).mean()
+
+    return test_loss, accuracy
+
+
+def train(model, trainloader, validloader, epochs, print_every, criterion, optimizer, device='cuda'):
+    epochs = epochs
+    print_every = print_every
+    steps = 0
+
+    # Change to train mode if not already
+    model.train()
+    # change to cuda
+    model.to(device)
+
+    for e in range(epochs):
+        running_loss = 0
+        for (images, labels) in trainloader:
+            steps += 1
+
+            images, labels = images.to(device), labels.to(device)
+
+            optimizer.zero_grad()
+
+            # Forward and backward passes
+            outputs = model.forward(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item()
+
+            if steps % print_every == 0:
+
+                # Make sure network is in eval mode for inference
+                model.eval()
+
+                # Turn off gradients for validation, saves memory and computations
+                with torch.no_grad():
+                    validation_loss, accuracy = validation(model, validloader, criterion)
+
+                print("Epoch: {}/{}.. ".format(e+1, epochs),
+                      "Training Loss: {:.3f}.. ".format(running_loss/print_every),
+                      "Validation Loss: {:.3f}.. ".format(validation_loss/len(validloader)),
+                      "Validation Accuracy: {:.3f}".format((accuracy/len(validloader))*100))
+
+                model.train()
+
+                running_loss = 0
+
+def check_accuracy_on_test(testloader, model):    
+    correct = 0
+    total = 0
+    model.to('cpu')
+    with torch.no_grad():
+        for data in testloader:
+            images, labels = data
+            outputs = model(images)
+            _, predicted = torch.max(outputs.data, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+    return 100 * correct / total
+
+def save_checkpoint(state, filename='checkpoint_extrainfo.pth'):
+    torch.save(state, filename)