FCN中单张图片可视化的代码

仅参考

import numpy as np
from PIL import Image
import sys
import matplotlib.pyplot as plt
#import scipy.io as sio
import caffe

#import caffe
import glob
#fileList=glob.glob('/home/yangwb/JPEGImages/*.jpg')
#f=open("/home/yangwb/MYwork/FCN/fcn.berkeleyvision.org-master/data/VOC2012/ImageSets/Segmentation/seg11valid1.txt")
#lines=f.readlines()
def vis_square(data):

    # 输入的数据为一个ndarray，尺寸可以为(n, height, width)或者是 (n, height, width, 3)
    # 前者即为n个灰度图像的数据，后者为n个rgb图像的数据
    # 在一个sqrt(n) by sqrt(n)的格子中，显示每一幅图像

    # 对输入的图像进行normlization
    data = (data - data.min()) / (data.max() - data.min())

    # 强制性地使输入的图像个数为平方数，不足平方数时，手动添加几幅
    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)   

    # 将所有输入的data图像平复在一个ndarray-data中（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的一个小例子,e.g., (3,120,120)
    # 即，这里的data是一个2d 或者 3d 的ndarray
    data = data.reshape((n * data.shape[1], n * data.shape[3]) + data.shape[4:])
    #print data
    # 显示data所对应的图像
    plt.imshow(data); plt.axis('off')
    plt.savefig('pool5.png')
#for i in range(1):
    #line = line.strip('\n')
    fileList='/home/yangwb/MYwork/FCN/fcn.berkeleyvision.org-master/myvoc1/2008_003821.jpg'
    #print fileList
    # load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
    #print filename
    im = Image.open(fileList)
    in_ = np.array(im, dtype=np.float32)
    in_ = in_[:,:,::-1]
    in_ -= np.array((104.00698793,116.66876762,122.67891434))
    in_ = in_.transpose((2,0,1))
    # init
    #caffe.set_device(int(sys.argv[1]))
    caffe.set_device(5)
    caffe.set_mode_gpu()
    # load net
    net = caffe.Net('/home/yangwb/MYwork/FCN/fcn.berkeleyvision.org-master/myvoc1/deploy.prototxt', '/home/yangwb/MYwork/FCN/fcn.berkeleyvision.org-master/myvoc1/_iter_45000.caffemodel', caffe.TEST)
    # shape for input (data blob is N x C x H x W), set data
    net.blobs['data'].reshape(1, *in_.shape)
    print net.blobs['data'].shape
    net.blobs['data'].data[...] = in_
    # run net and take argmax for prediction
    net.forward()
    out = net.blobs['score'].data[0].argmax(axis=0)
    #out=net.blobs['score'].data[0].argmax(0).flatten()
    #a=net.blobs['score'].data[0, 0].flatten()
    print out
    #abc='/home/yangwb/test1/'+line+'.jpg'
    #np.save(abc[:-4], out)
    score= net.blobs['upscore8'].data[0,15:16]
    #score= net.blobs['upscore21'].data[0]
    #print score[20]
    vis_square(score)

    #plt.imshow(out,cmap='gray')
    #plt.axis('off')
    #plt.savefig(fileList[:-4]+'.png')
    #plt.show()