import json
import torch
from torchvision import datasets, transforms, models

def process_image(image):...

def loading_data(data_dir):

	train_dir = '../training_dir/' + data_dir + '/train'
	valid_dir = '../training_dir/' + data_dir + '/valid'
	test_dir = '../training_dir/' + data_dir + '/test'

	train_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])])

	test_transforms = transforms.Compose([transforms.Resize(255),
                                      transforms.CenterCrop(224),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.485, 0.456, 0.406],
                                                           [0.229, 0.224, 0.225])])

	train_data = datasets.ImageFolder(train_dir, transform=train_transforms)
	test_data = datasets.ImageFolder(test_dir, transform=test_transforms)
	valid_data = datasets.ImageFolder(valid_dir, transform=test_transforms)

	trainloader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
	testloader = torch.utils.data.DataLoader(test_data, batch_size=64)
	validloader = torch.utils.data.DataLoader(valid_data, batch_size=64)

	return train_data, test_data, valid_data, trainloader, testloader, validloader

def extract_mapping(mapping_file_path, classes):...

import torch
from torch import nn
import torch.nn.functional as F
from torch import optim
from torchvision import datasets, transforms, models
from PIL import Image
import numpy as np
import data_utils as du

# Set the model to use
def set_pretrained_model(model_name):

def set_model_classifier(model, hidden_layer, input_size=25088, output_size=102, dropout=0.5):

# Train
def train_model(model, trainloader, set_epochs, validloader, learning_rate, device):

# Save the checkpoint
def create_checkpoint(model, model_name, train_data, epochs, optimizer, checkpoint_file_path, input_size=25088, output_size=102):

# Load a model from a checkpoint file
def load_checkpoint(checkpoint_file_path):

# Predict the class from an image file
def predict(image_path, checkpoint, device, topk=5):

import argparse
import torch
import os
import sys
sys.path.append('../utils_dir/')
import data_utils as du
import model_utils as mu

def main():

    input_arguments = process_arguments()
    default_device = torch.device("cuda" if torch.cuda.is_available() and input_arguments.gpu else "cpu")
    input_size = 25088
    chosen_architecture = input_arguments.chosen_archi
    if (chosen_architecture == "densenet121"): input_size = 1024
    if (chosen_architecture != "vgg16") and (chosen_architecture != "densenet121"):
        print("Pretrained Chosen architecture is densenet121 or vgg16, using default: vgg16")
        chosen_architecture = "vgg16"
        
    # Load data
    train_data, test_data, valid_data, trainloader, testloader, validloader = du.loading_data(input_arguments.data_directory)
    model = mu.set_pretrained_model(chosen_architecture)
    model = mu.set_model_classifier(model, input_arguments.hidden_units, input_size, output_size=102, dropout=0.05)
    
    # Train
    model, epochs, optimizer = mu.train_model(model,...
    
    if not os.path.exists(input_arguments.save_directory):
        os.makedirs(input_arguments.save_directory)
    
    # path to store model for use in predictions
    checkpoint_file_path = os.path.join(input_arguments.save_directory, "model.pth")

    # Store the trained model
    mu.create_checkpoint(model, 
                         input_arguments.chosen_archi,
                         train_data, 
                         epochs, 
                         optimizer, 
                         checkpoint_file_path,
                         input_size,
                         output_size=102)
    
    pass

from flask import Flask
from flask import request
import urllib.request
import argparse
from argparse import Namespace
import torch
import sys
sys.path.append('../utils_dir/')
import data_utils as du
import model_utils as mu

def predict():

    input_arguments = Namespace(cat_name_file='cat_to_name.json', checkpoint_file_path='../model_dir/model.pth', gpu=False, input_image_path='img/1.jpg', topk=3)
    default_device = torch.device("cuda" if torch.cuda.is_available() and input_arguments.gpu else "cpu")
    
    # Predict
    probs, classes = mu.predict(input_arguments.input_image_path, 
                                input_arguments.checkpoint_file_path,
                                default_device,
                                input_arguments.topk)
    
    # Extract
    image_classifications = du.extract_mapping(input_arguments.cat_name_file, classes)
    return "Image classification is : " + image_classifications[0] + ", predicted with probability: " + str(probs[0])

app = Flask(__name__)
@app.route('/')
def flask_server():
    # Download and store image
    image = request.args.get('image')
    urllib.request.urlretrieve(image, "img/1.jpg")

    # return data
    return '<h1>'+predict()+'</h1><br><img src="'+image+'" />'