import sys
import os
#suppresses non warning or error output to the terminal, must be done before import caffe
os.environ['GLOG_minloglevel'] = '2' 
import caffe
import numpy as np
import matplotlib.pyplot as plt

model_def = "/home/maxlotz/Thesis/prototxts/fcn_deploy.prototxt"
model_weights = "/data2/maxlotz/caffe_models/nyud-fcn32s-color-heavy.caffemodel"

caffe.set_mode_cpu()

net = caffe.Net(model_def,      # defines the structure of the model
                model_weights,  # contains the trained weights
                caffe.TEST)     # use test mode (e.g., don't perform dropout)

def vis_square(data):
    """Take an array of shape (n, height, width) or (n, height, width, 3)
       and visualize each (height, width) thing in a grid of size approx. sqrt(n) by sqrt(n)"""
    
    # normalize data for display
    data = (data - data.min()) / (data.max() - data.min())
    
    # force the number of filters to be square
    n = int(np.ceil(np.sqrt(data.shape[0])))
    padding = (((0, n ** 2 - data.shape[0]),
               (0, 1), (0, 1))                 # add some space between filters
               + ((0, 0),) * (data.ndim - 3))  # don't pad the last dimension (if there is one)
    data = np.pad(data, padding, mode='constant', constant_values=1)  # pad with ones (white)
    
    # tile the filters into an image
    data = data.reshape((n, n) + data.shape[1:]).transpose((0, 2, 1, 3) + tuple(range(4, data.ndim + 1)))
    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])
    
    plt.imshow(data); plt.axis('off')
    plt.show()

filters = net.params['conv1_1'][0].data
vis_square(filters.transpose(0, 2, 3, 1))