ml-inory / crowdnet Goto Github PK

import sys
sys.path.insert(0, '/home/tumh/caffe/python')
import caffe
import cv2

caffe.set_mode_gpu()
caffe.set_device(0)

net = caffe.Net('/home/tumh/CrowdNet/prototxt/CrowdNet_deploy.prototxt', '/home/tumh/CrowdNet/model/CrowdNet_50000.caffemodel', caffe.TEST)

#  im = cv2.imread('/home/tumh/img/000104.jpg')
im = cv2.imread('/home/tumh/CrowdNet/sample/resize_256x256.jpg').astype(np.float32)
im = im.transpose((2,0,1))
im = im[None,:]
net.blobs['data'].reshape(*(im.shape))
out = net.forward(data=im)
dmap = net.blobs['dmap'].data
dmap = dmap.squeeze()
print dmap.shape
print out['sum']