test.py 文件源码

def gen_net(num):
    caffe.set_device(0)
    caffe.set_mode_gpu()

    filename = '2007_000032.jpg'
    im = Image.open(filename)
    m = np.asarray(im, dtype=np.float32)
    m = m[:,:,::-1]
    m -= np.array((104.00698793,116.66876762,122.67891434))
    m = m.transpose((2, 0, 1))

    net = caffe.Net(
        "train_val.prototxt",
        "train_iter_" + str(num) + ".caffemodel",
        # "/data/VGG16/caffemodel",
        # "../fcn-32s/good.caffemodel",
        caffe.TRAIN)

    net.blobs["data"].reshape(1, *m.shape)
    net.blobs["data"].data[...] = m
    net.forward()
    return net