test.py

#coding=utf-8
import os
import sys
import numpy as np

from PIL import Image

import net

def main():
    # location of depth module, config and parameters
    module_fn = 'models/depth.py'
    config_fn = 'models/depth.conf'#网络结构
    params_dir = 'weights/depth'#网络相关参数

    # load depth network
    machine = net.create_machine(module_fn, config_fn, params_dir)

    # demo image
    rgb = Image.open('demo_nyud_rgb.jpg')
    rgb = rgb.resize((320, 240), Image.BICUBIC)

    # build depth inference function and run
    rgb_imgs = np.asarray(rgb).reshape((1, 240, 320, 3))
    pred_depths = machine.infer_depth(rgb_imgs)

    # save prediction
    (m, M) = (pred_depths.min(), pred_depths.max())
    depth_img_np = (pred_depths[0] - m) / (M - m)
    depth_img = Image.fromarray((255*depth_img_np).astype(np.uint8))
    depth_img.save('demo_nyud_depth_prediction.png')


if __name__ == '__main__':
    main()