from bigdl.dllib.nncontext import *
sc = init_nncontext("Tfnet Example")
import sys
import tensorflow as tf
sys.path
slim_path = "/path/to/yourdownload/models/research/slim" # Please set this to the directory where you clone the tensorflow models repository
sys.path.append(slim_path)

"""
Explanation: Image classification using tensorflow pre-trained model
Tensorflow-Slim image classification model library provides both the implementation and pre-trianed checkpoint many popular convolution neural nets for image classification.
Using TFNet in Analytics-Zoo, we can easily load these pre-trained model and make distributed inference with only a few lines of code.
Add the slim image classification model library to $PYTHONPATH
End of explanation
"""


from nets import inception
slim = tf.contrib.slim

images = tf.placeholder(dtype=tf.float32, shape=(None, 224, 224, 3))

with slim.arg_scope(inception.inception_v1_arg_scope()):
    logits, end_points = inception.inception_v1(images, num_classes=1001, is_training=False)

sess = tf.Session()
saver = tf.train.Saver()
saver.restore(sess, "/path/to/yourdownload/checkpoint/inception_v1.ckpt") # You need to edit this path to the checkpoint you downloaded

"""
Explanation: Construct the inference graph and restore the checkpoint
End of explanation
"""


from bigdl.dllib.utils.tf import export_tf
export_tf(sess, "/path/to/yourdownload/models/tfnet/", inputs=[images], outputs=[logits])

"""
Explanation: Export the graph as a frozen inference graph
The export_tf utility function will frozen the tensorflow graph, strip unused operation according to the inputs and outputs and save it to the specified directory along with the input/output tensor names.
End of explanation
"""


from bigdl.orca.tfpark import TFNet
model = TFNet.from_export_folder("/path/to/yourdownload/models/tfnet")

"""
Explanation: Load to Analytics-Zoo
End of explanation
"""


import cv2
import numpy as np
import json
im = cv2.imread("test.jpg")
im = cv2.resize(im, (224, 224))
im = (im - 127.0) / 128.0
im = np.expand_dims(im, 0)

import json
with open("imagenet_class_index.json") as f:
    class_idx = json.load(f)
idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]

import numpy as np
result = model.predict(im).first()
print(idx2label[np.argmax(result, 0)])

"""
Explanation: Test it on one image
End of explanation
"""


tf.reset_default_graph()# if you want to test your code, you can use it to reset your graph and to avoid some mistakes

images = tf.placeholder(dtype=tf.float32, shape=(None, 224, 224, 3))

with slim.arg_scope(inception.inception_v1_arg_scope()):
    logits, end_points = inception.inception_v1(images, num_classes=1001)

sess = tf.Session()
saver = tf.train.Saver()
saver.restore(sess, "/path/to/yourdownload/checkpoint/inception_v1.ckpt")# You need to edit this path to the checkpoint you downloaded

"""
Explanation: Fine tune on the data
Construct the inference graph and restore the checkpoint
End of explanation
"""


from bigdl.dllib.utils.tf import export_tf
avg_pool = end_points['AvgPool_0a_7x7']
export_tf(sess, "/path/to/yourdownload/models/tfnet/", inputs=[images], outputs=[avg_pool])

"""
Explanation: Remove the dense layer of the graph(Inception_v1) and export the graph left as a frozen inference graph
The export_tf utility function will frozen the tensorflow graph, strip unused operation according to the inputs and outputs and save it to the specified directory along with the input/output tensor names. In this example, we use the "AvgPool_0a_7x7" as our end_points, freeze the graph accoring to "AvgPool_0a_7x7" and export this graph to the specific directory
End of explanation
"""


from bigdl.orca.tfpark import TFNet
amodel = TFNet.from_export_folder("/path/to/yourdownload/models/tfnet/")

"""
Explanation: Load to Analytics-Zoo
Load that freezed graph from the specific directory above
End of explanation
"""


from bigdl.dllib.nn.layer import Sequential,Transpose,Contiguous,Linear,ReLU, SoftMax, Reshape, View, MulConstant, SpatialAveragePooling
full_model = Sequential()
full_model.add(Transpose([(2,4), (2,3)]))
scalar = 1. /255
full_model.add(MulConstant(scalar))
full_model.add(Contiguous())
full_model.add(amodel)
full_model.add(View([1024]))
full_model.add(Linear(1024,2))

"""
Explanation: Add the layers you want to the model
Use sequential model to build your model. First transport the type of data from NCHW to NHWC. Then multiply the scalar and make the input and output both contiguous. Add the freezed graph and add a linear layer after the graph.
End of explanation
"""


import re
from bigdl.dllib.nn.criterion import CrossEntropyCriterion
from pyspark import SparkConf
from pyspark.ml import Pipeline
from pyspark.sql.functions import col, udf
from pyspark.sql.types import DoubleType, StringType
from bigdl.dllib.nncontext import *
from bigdl.dllib.feature.image import *
from bigdl.dllib.keras.layers import Dense, Input, Flatten
from bigdl.dllib.keras.models import *
from bigdl.dllib.net import *
from bigdl.dllib.nnframes import *
image_path = "file:///path/toyourdownload/dogs-vs-cats/train/*.jpg"
imageDF = NNImageReader.readImages(image_path, sc)

getName = udf(lambda row:
                  row[0],
                  StringType())
def label(name):
    if 'cat' in name:
        return 1.0
    elif 'dog' in name:
        return 2.0
    else:
        raise ValueError("file name format is not correct: %s" % name)
getLabel = udf(lambda name: label(name), DoubleType())

labelDF = imageDF.withColumn("name", getName(col("image"))) \
        .withColumn("label", getLabel(col('name')))
(trainingDF, validationDF) = labelDF.randomSplit([0.9, 0.1])
labelDF.select("name","label").show(10, False)

"""
Explanation: get data
Set the path to the cat_dog data you downloaded. Then read these images from this path.
Use udf functions to design some functions to mark data.
Seperate the data into two groups: training data and validation data.
End of explanation
"""


trainingDF.groupBy("label").count().show()

"""
Explanation: Show the distribution of Training Data
End of explanation
"""


transformer = ChainedPreprocessing(
        [RowToImageFeature(), ImageResize(224, 224),
         ImageChannelNormalize(123.0, 117.0, 104.0), ImageMatToTensor(), ImageFeatureToTensor()])

from bigdl.dllib.optim.optimizer import *
from bigdl.dllib.nn.criterion import *
import datetime as dt
app_name='classification cat vs dog'+dt.datetime.now().strftime("%Y%m%d-%H%M%S")
train_summary = TrainSummary(log_dir='./log/',
                                     app_name=app_name)
classifier = NNClassifier(full_model, CrossEntropyCriterion(), transformer)\
        .setFeaturesCol("image")\
        .setCachingSample(False)\
        .setLearningRate(0.003)\
        .setBatchSize(16)\
        .setMaxEpoch(9)\
        .setTrainSummary(train_summary)
# BatchSize is a multiple of physcial_core_number(in your bash command)

from pyspark.ml import Pipeline
import os
import datetime as dt
if not os.path.exists("./log"):
    os.makedirs("./log")
    
print("Saving logs to ", app_name)
pipeline = Pipeline(stages=[classifier])
trainedModel = pipeline.fit(trainingDF)

"""
Explanation: Using pipeline to train the data
Combine the preprocessing processes together. Before identify the classifier, design the summary's name to make sure we will store the log in right name. Identify the args of classifier. Then save logs to the dirctionary. Use the classifier identified just to create pipeline. Use this pipeline as your model to train the data.
End of explanation
"""


# predict using the validation data
predictionDF = trainedModel.transform(validationDF).cache()
# caculate the correct rate and the test error
correct = predictionDF.filter("label=prediction").count()
overall = predictionDF.count()
accuracy = correct * 1.0 / overall

print("Test Error = %g " % (1.0 - accuracy))

"""
Explanation: Test the model on validation data
Use the vallidation data to test the model, and print the test error
End of explanation
"""


# get a pic about the result and make it can show in GUI
import matplotlib
matplotlib.use('Agg')
%pylab inline
from matplotlib import pyplot as plt
from matplotlib.pyplot import imshow

"""
Explanation: Make the pic can be seen on this page. Repeatedly run this code without restart the kernel will result in a warning.
End of explanation
"""


# read loss from summary nad show in pic
loss = np.array(train_summary.read_scalar("Loss"))
plt.figure(figsize = (12,12))
plt.plot(loss[:,0], loss[:,1], label='loss')
plt.xlim(0, loss.shape[0]+10)
plt.grid(True)
plt.title("loss")

"""
Explanation: Show the change in loss using a pic.
End of explanation
"""


samplecat=predictionDF.filter(predictionDF.prediction==1.0).limit(2).collect()
sampledog=predictionDF.filter(predictionDF.prediction==2.0).sort("label", ascending=False).limit(2).collect()

"""
Explanation: Collect two examples from both sides.
End of explanation
"""


from IPython.display import Image, display
for cat in samplecat:
    print ("prediction:"), cat.prediction
    display(Image(cat.image.origin[5:], height=256,width=256))

"""
Explanation: Show two cat-pics and their predictions.
End of explanation
"""


# show the two pics and their things as follow
from IPython.display import Image, display
for dog in sampledog:
    print ("prediction:"), dog.prediction
    
    display(Image(dog.image.origin[5:], height=256,width=256))

"""
Explanation: Show two dog-pics and their predictions.
End of explanation
"""


