# ----------------------------------------------------------------------------

# Copyright 2016 Nervana Systems Inc.

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ----------------------------------------------------------------------------

import os

import sys

import h5py

import pickle

import logging

import argparse

import numpy as np

from neon import NervanaObject

from neon.backends import gen_backend

from neon.models import Model

from neon.data import ImageParams

from neon.models import Model

from neon.util.persist import load_obj

from neon.util.modeldesc import ModelDescription

from segnet_neon_backend import NervanaGPU_Upsample

from segnet_neon_backend import _get_bprop_upsampling

from pixelwise_dataloader import PixelWiseImageLoader

from segnet_neon import gen_model

global _get_bprop_upsampling

parser = argparse.ArgumentParser(description='compare output of trained model to ground truth')

parser.add_argument('image_path', type=str, help='path to the CamVid data set in the SegNet-Tutorial '

'GitHub repo')

parser.add_argument('save_model_file', type=str, help='serialized model file')

parser.add_argument('output_dir', type=str, help='directory to save the output images')

parser.add_argument("--num_classes", type=int, default=12, help="number of annotated classes")

parser.add_argument("--height", type=int, default=256, help="image height")

parser.add_argument("--width", type=int, default=512, help="image width")

parser.add_argument("--display", action="store_true", help="output to screen instead of to file")

args = parser.parse_args()

if not args.display:

print 'Will try to output to the display'

import matplotlib

matplotlib.use('Agg')

import matplotlib.pyplot as plt

if not os.path.isdir(args.output_dir):

print 'making output diriextory [%s]' % args.output_dir

os.makedirs(args.output_dir)

else:

print 'will put output images in directory %s' % args.output_dir

# check that image dimensions are powers of 2

if((args.height & (args.height - 1)) != 0):

raise TypeError("Height must be a power of 2.")

if((args.width & (args.width - 1)) != 0):

raise TypeError("Width must be a power of 2.")

(c, h, w) = (args.num_classes, args.height, args.width)

# need to use the backend with the new upsampling layer implementation

be = NervanaGPU_Upsample(device_id=0)

# set batch size

be.bsz = 1

# couple backend to global neon object

NervanaObject.be = be

# input image parameters - channel_count is 3 for color images, 1 for B/W

shape = dict(channel_count=3, height=h, width=w, subtract_mean=False)

test_params = ImageParams(center=True, flip=False,

scale_min=min(h, w), scale_max=min(h, w),

aspect_ratio=0, **shape)

common = dict(target_size=h*w, target_conversion='read_contents',

onehot=False, target_dtype=np.uint8, nclasses=c)

data_dir = args.image_path

test_set = PixelWiseImageLoader(set_name='test', repo_dir=data_dir,media_params=test_params,

index_file=os.path.join(data_dir, 'test_images.csv'), **common)

# initialize model object

segnet_model = Model(layers=gen_model(c, h, w))

segnet_model.initialize(test_set)

# load up the serialized model

model_desc = ModelDescription(load_obj(args.save_model_file))

for layer in segnet_model.layers_to_optimize:

name = layer.name

trained_layer = model_desc.getlayer(name)

layer.load_weights(trained_layer)

fig = plt.figure()

if args.display:

plt.ion()

im1 = None

im2 = None

cnt = 1

for x, t in test_set:

z = segnet_model.fprop(x).get()

z = np.argmax(z.reshape((c, h, w)), axis=0)

t = np.argmax(t.get().reshape((c, h, w)), axis=0)

# calculate the misclass rate

acc = (np.where(z == t)[0].size / float(z.size))*100.0

plt.subplot(2,1,1);

if im1 is None:

im1 = plt.imshow(t);plt.title('Truth')

else:

im1.set_data(t)

plt.subplot(2,1,2);

if im2 is None:

im2 = plt.imshow(z);

else:

im2.set_data(z)

plt.title('SegNet - pixelwise acc = %.1f %%' % acc )

fig.savefig(os.path.join(args.output_dir, 'out_%d.png' % cnt))

cnt += 1

if args.display:

plt.draw()

plt.show()

raw_input('Pixelwise classification accuracy = %.1f, press return to continue...' % acc)