How to feed HDF5 testing data into a trained Unet CNN by pycaffe

[Tianyi Miao]

Hi all,

I am new in caffe, and learning on U Net. Now I am able to train my own network by using their prototxt, training set is stored in HDF5 file. My question is how should use this trained network to produce result by using pycaffe, following is what I have tried so far:

from PIL import Image
import numpy as np
import caffe 
import h5py
NET_FILE = '/.../.../deploy.prototxt'
MODEL_FILE = '/.../.../my.caffemodel'
net = caffe.Net(NET_FILE, MODEL_FILE, caffe.TEST)

with h5py.File('/.../.../forPrediction.h5', 'r') as f:
   firstbatch = f['data'][0:1,:,:,:]       #(1,1,512,512) grayscale image

net.blobs['data'].data[...]=firstbatch

out = net.forward()                        #(1,1,340,340) predict label

imageChannel0=Image.fromarray(out['score'][0,0,:,:]*255, 'L')     #predict range (0-1)
imageChannel1=Image.fromarray(out['score'][0,1,:,:]*255, 'L')     #rescale to (0-255)

imageChannel0.save('/.../.../foreground.png')
imageChannel1.save('/.../.../background.png')

Is that a correct process?

I asking that is because when I check the output two images, they looks wried, as I attached. So I really want to figure out is that because of the mistake from NN training or from pycaffe prediction.

And BTW if someone familiar with U-NET, could I ask why the final layer has 2 layers(foreground and background as they introduce)? how should I prepare label in HDF5 file(e.g.'f['label'][0,0,:,:]=foreground,f['label'][0,1,:,:]=background')

Any help appreciated!