import os
import errno


_base_model_url = 'http://data.mxnet.io/models/'
_default_model_info = {
    'imagenet11k-resnet-152': {'symbol':_base_model_url+'imagenet-11k/resnet-152/resnet-152-symbol.json',
                             'params':_base_model_url+'imagenet-11k/resnet-152/resnet-152-0000.params'},
}


def download_file(url, local_fname=None, force_write=False):
    # requests is not default installed
    import requests
    if local_fname is None:
        local_fname = url.split('/')[-1]
    if not force_write and os.path.exists(local_fname):
        return local_fname

    dir_name = os.path.dirname(local_fname)

    if dir_name != "":
        if not os.path.exists(dir_name):
            try:  # try to create the directory if it doesn't exists
                os.makedirs(dir_name)
            except OSError as exc:
                if exc.errno != errno.EEXIST:
                    raise

    r = requests.get(url, stream=True)
    assert r.status_code == 200, "failed to open %s" % url
    with open(local_fname, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:  # filter out keep-alive new chunks
                f.write(chunk)
    return local_fname


def download_model(model_name, dst_dir='./', meta_info=None):
    if meta_info is None:
        meta_info = _default_model_info
    meta_info = dict(meta_info)
    if model_name not in meta_info:
        return (None, 0)
    if not os.path.isdir(dst_dir):
        os.mkdir(dst_dir)
    meta = dict(meta_info[model_name])
    assert 'symbol' in meta, "missing symbol url"
    model_name = os.path.join(dst_dir, model_name)
    download_file(meta['symbol'], model_name+'-symbol.json')
    assert 'params' in meta, "mssing parameter file url"
    download_file(meta['params'], model_name+'-0000.params')
    return (model_name, 0)


if __name__ == "__main__":
    # ***** Download synset (i.e., Synonym Set):
    synset_url = 'http://data.mxnet.io.s3-website-us-west-1.amazonaws.com/models/imagenet-11k/synset.txt'
    download_file(synset_url, 'synset.txt')

    # ***** Download Model:
    download_model('imagenet11k-resnet-152', dst_dir='./')

import torch
import numpy as np
from tensorflow.contrib.keras.api.keras.preprocessing import image


# ************** Parameters:
num_predictions = 5  # Top-k Results
model_address = 'resnet152Full.pth'  # for loading models
lexicon_address = 'synset.txt'
test_image_address = 'seagull.jpg'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# Load Converted Model:
model = torch.load(model_address).to(device)
model.eval()


# Read Input Image and Apply Pre-process:
img = image.load_img(test_image_address, target_size=(224, 224))
x = image.img_to_array(img)
x = x[..., ::-1]  # transform image from RGB to BGR
x = np.transpose(x, (2, 0, 1))
x = np.expand_dims(x, 0).copy()
x = torch.from_numpy(x)
x = x.to(device)


# Load Full-ImageNet Dictionary (i.e., lexicon):
with open(lexicon_address, 'r') as f:
    labels = [l.rstrip() for l in f]


# Make prediction (forward pass):
with torch.no_grad():
    output = model(x)
max, argmax = output.data.squeeze().max(0)
class_id = argmax.item()
class_name = labels[class_id]


# Print the top-5 Results:
h_x = output.data.squeeze()
probs, idx = h_x.sort(0, True)
print('Top-5 Results: ')
for i in range(0, num_predictions):
    print('{:.2f}% -> {}'.format(probs[i] * 100.0, labels[idx[i]]))
str_final_label = 'The Image is a ' + class_name[10:] + '.'
print(str_final_label)