Final

predict.py

@@ -0,0 +1,158 @@
+import os
+import argparse
+import torch
+import json
+from collections import OrderedDict
+from torchvision import models
+from torch import nn
+import numpy as np
+from PIL import Image
+
+
+
+
+def process_image(image):
+    ''' Scales, crops, and normalizes a PIL image for a PyTorch model,
+        returns an Numpy array
+    '''
+    np_img = None
+    with Image.open(image) as im:
+        w, h = im.size
+        min_s = min(w, h)
+        if min_s == w:
+            w = 256
+            h = h * 256 // w
+        else:
+            h = 256
+            w = w * 256 // h
+        im.thumbnail((w, h))
+        w, h = im.size
+        (left, upper, right, lower) = w//2-224//2, h//2-224/2, w//2+224//2, h//2+224//2
+        im_cropped = im.crop((left, upper, right, lower))
+
+        np_img = np.array(im_cropped) / 255
+    arr = (np_img - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])
+    return torch.from_numpy(arr.transpose(2, 0, 1))
+
+
+def predict(image_path, model, device, topk=5):
+    ''' Predict the class (or classes) of an image using a trained deep learning model.
+    '''
+    probs, classes = None, None
+    image = process_image(image_path)
+    image = image.view((1, 3, 224, 224)).type(torch.FloatTensor)
+
+    model.eval()
+    model.to(device)
+    with torch.no_grad():
+
+        image = image.to(device)
+        output = model(image)
+        ps = torch.exp(output)
+
+        top_p, top_class = ps.topk(topk, dim=1)
+        probs, classes = top_p.tolist()[0], top_class.tolist()[0]
+    return probs, classes
+
+
+def create_arguments(parser):
+    parser.add_argument("img_pth", type=str)
+    parser.add_argument("checkpoint_pth", type=str)
+    parser.add_argument("--top_k", type=int,default=1)
+    parser.add_argument("--category_names",type=str,default="cat_to_name.json")
+    parser.add_argument("--gpu", action='store_true')
+
+
+def check_data(category_names, img_pth, checkpoint_pth, top_k):
+    with open(category_names, 'r') as f:
+        cat_to_name = json.load(f, strict=False)
+    if not os.path.exists(img_pth):
+        raise ValueError("Image Path not exists")
+    if not os.path.exists(checkpoint_pth):
+        raise ValueError("Checkpoint Path not exists")
+    if top_k <= 0:
+        raise ValueError()
+    return cat_to_name
+
+
+def restore_model(checkpoint_pth):
+    # Load model
+    checkpoint = torch.load(checkpoint_pth)
+    model_name = checkpoint["model"]
+    class_to_idx = checkpoint["class_to_idx"]
+    hidden_units = checkpoint["hidden_units"]
+    state_dict = checkpoint["state_dict"]
+
+    model = getattr(models, model_name)(pretrained=True)
+    for param in model.parameters():
+        param.requires_grad = False
+
+    if model.__class__.__name__ == "ResNet":
+        # ResNet classify layer is fc
+        in_features = model.fc.in_features
+
+        fc = nn.Sequential(OrderedDict([
+            ("fc1", nn.Linear(in_features, hidden_units)),
+            ("relu1", nn.ReLU()),
+            ("fc2", nn.Linear(hidden_units, 102)),
+            ("output", nn.LogSoftmax(dim=1))
+        ]))
+
+        model.fc = fc
+    elif model.__class__.__name__ == "VGG":
+        # VGG classify layer is classifer
+        # It has 6 mini-layers
+        classifier = nn.Sequential(OrderedDict([
+            ("fc1", nn.Linear(25088, 4096)),
+            ("relu1", nn.ReLU()),
+            ("dropout1", nn.Dropout(0.5)),
+            ("fc2", nn.Linear(4096, hidden_units)),
+            ("relu2", nn.ReLU()),
+            ("dropout2", nn.Dropout(0.5)),
+            ("fc3", nn.Linear(hidden_units, 102)),
+            ("output", nn.LogSoftmax(dim=1))
+        ]))
+        model.classifier = classifier
+    model.load_state_dict(state_dict=state_dict)
+    model.class_to_idx = class_to_idx
+    
+    return model
+
+
+def show_result(cat_to_name, probs, classes, class_to_idx):
+    if cat_to_name:
+        res = dict((v, k) for k, v in class_to_idx.items())
+        types = [cat_to_name[res[i]] for i in classes]
+
+        idx = probs.index(max(probs))
+        print(f"Class {types[idx]} with prob: {probs[idx]}")
+
+        print(f"Top {top_k}")
+        for i in range(len(probs)):
+            print(f"Class {types[i]} with prob: {probs[i]}")
+    else:
+        idx = probs.index(max(probs))
+        print(f"Class {classes[idx]} with prob: {probs[idx]}")
+
+        print(f"Top {top_k}")
+        for i in range(len(probs)):
+            print(f"Class {classes[i]} with prob: {probs[i]}")
+
+
+if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser()
+    create_arguments(arg_parser)
+    args = arg_parser.parse_args()
+
+    img_pth = args.img_pth
+    checkpoint_pth = args.checkpoint_pth
+    top_k = args.top_k
+    gpu = args.gpu
+    cat_to_name = check_data(args.category_names, img_pth, checkpoint_pth, top_k)
+    device = torch.device('cuda' if torch.cuda.is_available() and args.gpu else 'cpu')
+    
+    
+    model = restore_model(checkpoint_pth)
+    probs, classes = predict(img_pth, model, device, top_k)
+    show_result(cat_to_name, probs, classes, model.class_to_idx)
+

train.py

@@ -0,0 +1,176 @@
+import argparse
+import torch
+from torch import nn, optim
+from torchvision import datasets, models, transforms
+from torch.utils.data.dataloader import DataLoader
+import os
+from collections import OrderedDict
+
+
+# Transforms
+train_transform = transforms.Compose([
+    transforms.Resize(255),
+    transforms.ColorJitter(brightness=2),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(30),
+    transforms.RandomResizedCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+valid_transform = transforms.Compose([
+    transforms.Resize(255),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+
+def create_arguments(parser):
+    parser.add_argument("data_directory", type=str)
+    parser.add_argument("--save_dir", type=str, help="Directory to save checkpoints", default=os.getcwd())
+    parser.add_argument("--arch", type=str, help="Model's architecture", default="resnet34")
+    parser.add_argument("--learning_rate", type=float, help="Model's learning rate", default=0.002)
+    parser.add_argument("--hidden_units", type=int, help="Model's number of hidden units", default=128)
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--gpu", action='store_true')
+
+
+def check_data(data_directory, save_dir, hidden_units, learning_rate, epochs):
+    # Check data
+    if not os.path.exists(data_directory):
+        raise ValueError("Data directory not exist!")
+    if not os.path.exists(save_dir):
+        raise ValueError("Checkpoint save directory not exist!")
+    if hidden_units <= 0 or learning_rate <= 0 or epochs <= 0:
+        raise ValueError()
+
+
+def create_model(arch, hidden_units, learning_rate):
+    model = getattr(models, arch)(pretrained=True)
+    optimizer = None
+    for param in model.parameters():
+        param.requires_grad = False
+
+    if model.__class__.__name__ == "ResNet":
+        # ResNet classify layer is fc
+        in_features = model.fc.in_features
+
+        fc = nn.Sequential(OrderedDict([
+            ("fc1", nn.Linear(in_features, hidden_units)),
+            ("relu1", nn.ReLU()),
+            ("fc2", nn.Linear(hidden_units, 102)),
+            ("output", nn.LogSoftmax(dim=1))
+        ]))
+
+        model.fc = fc
+
+        optimizer = optim.Adam(model.fc.parameters(), lr=learning_rate)
+    elif model.__class__.__name__ == "VGG":
+        # VGG classify layer is classifer
+        # It has 6 mini-layers
+        classifier = nn.Sequential(OrderedDict([
+            ("fc1", nn.Linear(25088, 4096)),
+            ("relu1", nn.ReLU()),
+            ("dropout1", nn.Dropout(0.5)),
+            ("fc2", nn.Linear(4096, hidden_units)),
+            ("relu2", nn.ReLU()),
+            ("dropout2", nn.Dropout(0.5)),
+            ("fc3", nn.Linear(hidden_units, 102)),
+            ("output", nn.LogSoftmax(dim=1))
+        ]))
+        model.classifier = classifier
+        optimizer = optim.SGD(model.classifier.parameters(), lr=learning_rate)
+    else:
+        raise ValueError("Architecture not support")
+    
+    return model, optimizer
+
+
+
+def training(model, optimizer, criterion, epochs, device, trainloader, validloader):
+    train_losses, valid_losses = [], []
+    for i in range(epochs):
+        print(f"Epoch {i + 1}/{epochs}")
+        model.train()
+
+        train_loss = 0
+        valid_loss = 0
+        accuracy = 0
+        for images, labels in trainloader:
+            images, labels = images.to(device), labels.to(device)
+            output = model(images)
+
+            loss = criterion(output, labels)
+            train_loss += loss.item()
+            
+            
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        train_losses.append(train_loss)
+        with torch.no_grad():
+            model.eval()
+            for images, labels in validloader:
+                images, labels = images.to(device), labels.to(device)
+                output = model(images)
+
+                loss = criterion(output, labels)
+                valid_loss += loss.item()
+
+                output_p = torch.exp(output)
+                _, top_class = output_p.topk(1, dim=1)
+                equals = top_class == labels.view(top_class.shape)
+                accuracy += torch.mean(equals.type(torch.FloatTensor))
+        accuracy = accuracy / len(validloader) * 100
+        
+        print(f"Train loss: {train_loss}")
+        print(f"Valid loss: {valid_loss}")
+        print(f"Valid accuracy: {accuracy}")
+        print("---------------------")
+
+def save_checkpoint(class_to_idx, arch, hidden_units, model, save_dir):
+    checkpoint = {
+        "class_to_idx": class_to_idx,
+        "model": arch,
+        "hidden_units": hidden_units,
+        "state_dict": model.state_dict()
+    }
+
+    torch.save(checkpoint, os.path.join(save_dir, "checkpoint.pth"))
+
+if __name__ == '__main__':
+    arg_parser = argparse.ArgumentParser()
+    create_arguments(arg_parser)
+    args = arg_parser.parse_args()
+
+    data_directory = args.data_directory
+    save_dir = args.save_dir
+    arch = args.arch
+    learning_rate = args.learning_rate
+    hidden_units = args.hidden_units
+    epochs = args.epochs
+    gpu = args.gpu
+    device = torch.device('cuda' if torch.cuda.is_available() and args.gpu else 'cpu')
+    check_data(data_directory, save_dir, hidden_units, learning_rate, epochs)
+    
+    train_data = os.path.join(data_directory, "train")
+    valid_data = os.path.join(data_directory, "valid")
+    train_dataset = datasets.ImageFolder(train_data, transform=train_transform)
+    trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
+
+    valid_dataset = datasets.ImageFolder(valid_data, transform=valid_transform)
+    validloader = DataLoader(valid_dataset, batch_size=64, shuffle=True)
+    
+    model, optimizer = create_model(arch, hidden_units, learning_rate)
+    
+    criterion = nn.NLLLoss()
+
+    model.to(device)
+    criterion.to(device)
+    
+    training(model, optimizer, criterion, epochs, device, trainloader, validloader)
+    save_checkpoint(train_dataset.class_to_idx, arch, hidden_units, model, save_dir)
+    
+