from __future__ import division, print_function


import sys

import python.caffe as caffe
import numpy as np
import os

DATA_DIR = "/home/sandbox/git/caffe/python/Visual_Recognition_HRA_v0_2/caffe2Keras/caffe_models"

OUTPUT_DIR = os.path.join(DATA_DIR, "vgg16_places365/saved_weights")


CAFFE_HOME="/home/sandbox/git/caffe/python/Visual_Recognition_HRA_v0_2/caffe2Keras/"

MODEL_DIR = os.path.join(CAFFE_HOME, "caffe_models", "vgg16_places365")
MODEL_PROTO = os.path.join(MODEL_DIR, "deploy.prototxt")
MODEL_WEIGHTS = os.path.join(MODEL_DIR, "vgg16_places365.caffemodel")
MEAN_IMAGE = os.path.join(MODEL_DIR, "places365CNN_mean.binaryproto")

caffe.set_mode_cpu()
net = caffe.Net(MODEL_PROTO, MODEL_WEIGHTS, caffe.TEST)


# layer names and output shapes
for layer_name, blob in net.blobs.iteritems():
    print(layer_name, blob.data.shape)

# write out weight matrices and bias vectors
for k, v in net.params.items():
    print(k, v[0].data.shape, v[1].data.shape)
    np.save(os.path.join(OUTPUT_DIR, "W_{:s}.npy".format(k)), v[0].data)
    np.save(os.path.join(OUTPUT_DIR, "b_{:s}.npy".format(k)), v[1].data)

# write out mean image
blob = caffe.proto.caffe_pb2.BlobProto()
with open(MEAN_IMAGE, 'rb') as fmean:
    mean_data = fmean.read()
blob.ParseFromString(mean_data)
mu = np.array(caffe.io.blobproto_to_array(blob))
print("Mean image:", mu.shape)
np.save(os.path.join(OUTPUT_DIR, "mean_image.npy"), mu)