predict

#! /usr/bin/env python

import json
import os
import sys
from PIL import Image
from torchvision import transforms
from lib.options import parse_predict
from lib.checkpoint import Checkpoint
from lib.config_parser import ConfigParser
from lib.logger import Logger

# Parse prediction options.
options, args = parse_predict()

# Expects two positional arguments.
if len(args) != 2:
  Logger.error("Missing required positional arguments: image_file and checkpoint.")
  sys.exit()

# Override defaults.
k = options.k or 3
device = 'cuda' if options.gpu else 'cpu'

# Load the saved model.
model = Checkpoint.load(args[1])

# Load image and transform.
image = Image.open(args[0])

# Define pre-processing transformations.
preprocess = transforms.Compose([
  transforms.Resize(225),
  transforms.CenterCrop(224),
  transforms.ToTensor(),
  transforms.Normalize(
    [0.485, 0.456, 0.406],
    [0.229, 0.224, 0.225]
  )
])

# Run transforms.
image = preprocess(image)

# Predict the flower classes.
classes, probs = model.predict(image, k=k, device=device)

if options.category_names:
  # Load class-to-name mapping.
  with open(options.category_names) as f:
    classes_to_names = json.load(f)

    # Convert to friendly names.
    names = [classes_to_names[cls] for cls in classes]

    # Create table data.
    data = [[n, p] for n, p in zip(names, probs)]
    headers = ['Name', 'Probability']
else:
  # Create table data.
  data = [[c, p] for c, p in zip(classes, probs)]
  headers = ['Class', 'Probability']

# Show results table.
Logger.table(data, headers=headers, color='m')