FCN语义分割初探——使用训练好的模型进行分割

前言

由于课题的需要，最近开始研究FCN语义分割，这几天将环境搭建好了立即测试了一下，这里分享出来与大家进行分享。

1. 准备

1.1 运行环境

这里使用到的环境是Ubuntu下PyCaffe，具体的环境搭建大家可以参考我的这篇博客进行环境搭架。Ubuntu16.04下安装Caffe记录（GPU）
这里使用到的FCN语义分割源码可以从Github上进行下载，fcn.berkeleyvision.org，本次实验使用到的是fcn8s-heavy-pascal.caffemodel。具体的网址的在该文件夹下的caffemodel-url中，复制里面的内容迅雷下载就好了，文件大概500多M

2. 运行

这次使用的到主要为infer.py这个文件，这里直接给出内容

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

import sys
sys.path.append('/home/sucker/Desktop/caffe-1.0/python')
sys.path.append('/home/sucker/Desktop/caffe-1.0/python/caffe')
sys.path.append('/home/sucker/Desktop/caffe-1.0/build/lib')
import caffe
#注意我这里没有将Pycaffe的目录导入到环境变量里面，直接添加的目录

# load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
im = Image.open('timg.jpeg')#网上找的测试图片
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))

# load net
net = caffe.Net('deploy.prototxt', 'fcn8s-heavy-pascal.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)
net.blobs['data'].data[...] = in_

# run net and take argmax for prediction
net.forward()
out = net.blobs['score'].data[0].argmax(axis=0)

plt.imshow(out,cmap='gray');
plt.axis('off')
plt.savefig('test1.png')

修改好之后，就开始进行测试，在infer.py目录下输入

python infer.py

运行完成之后就可以看到，分隔好的图像了

ref：caffe初步实践———使用训练好的模型完成语义分割任务