from __future__ import division

import caffe
import numpy as np
import cv2
import matplotlib


caffe.set_mode_gpu()
caffe.set_device(0)
net = caffe.Net('deploy.prototxt', 'blurMapping.caffemodel', caffe.TEST)

pixel_means= np.array([[[104.008, 116.669, 122.675]]])


img = cv2.imread('motion0001.jpg', 1)
img = img.astype(np.float32)
H = img.shape[0]
W = img.shape[1]
max_shape = max(img.shape[0], img.shape[1])
scale = 384.0/max_shape
img = img - pixel_means  
img = cv2.resize(img, None, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
h = img.shape[0]
w = img.shape[1]
img_square = np.zeros((384,384,3), dtype=np.float32)
img_square[0:img.shape[0],0:img.shape[1],:] = img
img = img_square.copy()
img = img[None, :, :, :]
img = img.transpose(0,3,1,2)
blobs = {}
blobs['data'] = img
net.blobs['data'].reshape(*(blobs['data'].shape))
forward_kwargs = {'data': blobs['data'].astype(np.float32, copy=False)}
blobs_out = net.forward(**forward_kwargs)
score = blobs_out['score']
score = score[0,0,:,:]
score = cv2.resize(score, (384, 384), interpolation=cv2.INTER_LINEAR)
score = score[0:h,0:w]
score = cv2.resize(score, (W, H), interpolation=cv2.INTER_NEAREST)
cv2.imwrite('soft_blur.png', np.uint8(255.0*score))
cv2.imwrite('hard_blur.png', np.uint8(255.0*(score > 0.5))) 
