fast_neural_style.py 文件源码

def main():
    parse_args()
    if (opt.input_image == '') and (opt.input_dir == ''):
        raise Exception('Must give exactly one of -input_image or -input_dir')
    checkpoint = load_t7checkpoint(opt.model, custom_layers=custom_layers)
    model = checkpoint.model
    model.evaluate()
    gc.collect()
    preprocess_method = checkpoint.opt.preprocessing or 'vgg'
    preprocess = methods[preprocess_method]

    def run_image(in_path, out_path):
        img = imread(in_path)
        img = np.array(img, dtype=np.float64)
        if opt.image_size > 0:
          img = scipy.misc.imresize(img,  np.float(opt.image_size)/np.float(np.max(img.shape))) #
#(768, 1153, 3)) # FIXME: IT WORKS ONLY WITH THESE DIMS
        import pdb; pdb.set_trace()
        img = img.transpose(2, 0, 1)
        _, H, W = img.shape
        img = img.reshape(1, 3, H, W)
        img_pre = preprocess.preprocess(img)
        img_out = model.forward(img_pre)

        img_out = preprocess.deprocess(img_out)[0]
        img_out = img_out.transpose(1, 2, 0)
        if opt.median_filter > 0:
            img_out = scipy.ndimage.filters.median_filter(
                img_out, opt.median_filter)
        scipy.misc.imsave(out_path, img_out)
        print('Writing output image to ' + out_path)
        outdir = os.path.dirname(out_path)
        if outdir is not '' and not os.path.exists(outdir):
            os.makedirs(outdir)
        scipy.misc.imsave(out_path, img_out)

    if opt.input_dir != '':
        if opt.output_dir == '':
            raise Exception('Must give -output_dir with -input_dir')
        for fn in os.path.isfile(opt.input_dir):
            if is_image_file(fn):
                in_path = os.path.concat(opt.input_dir, fn)
                out_path = os.path.concat(opt.output_dir, fn)
                run_image(in_path, out_path)

    elif opt.input_image != '':
        if opt.output_image == '':
            raise Exception('Must give -output_image with -input_image')
        run_image(opt.input_image, opt.output_image)