style_transfer.py 文件源码

def main(style_img, gpu_id, content_img, verbose, model, init, num_iters, ratio, length, output):
    """
        Entry point.
    """

    # logging
    level = logging.INFO if verbose else logging.DEBUG
    logging.basicConfig(format=LOG_FORMAT, datefmt="%H:%M:%S", level=level)
    logging.info("Starting style transfer.")

    # set GPU/CPU mode
    if gpu_id == -1:
        caffe.set_mode_cpu()
        logging.info("Running net on CPU.")
    else:
        caffe.set_device(gpu_id)
        caffe.set_mode_gpu()
        logging.info("Running net on GPU {0}.".format(gpu_id))

    # load images
    img_style = caffe.io.load_image(style_img)
    img_content = caffe.io.load_image(content_img)
    logging.info("Successfully loaded images.")

    # artistic style class
    use_pbar = not verbose
    st = StyleTransfer(model.lower(), use_pbar=use_pbar)
    logging.info("Successfully loaded model {0}.".format(model))

    # perform style transfer
    start = timeit.default_timer()
    n_iters = st.transfer_style(img_style, img_content, length=length,
                                init=init, ratio=np.float(ratio),
                                n_iter=num_iters, verbose=verbose)
    end = timeit.default_timer()
    logging.info("Ran {0} iterations in {1:.0f}s.".format(n_iters, end - start))
    img_out = st.get_generated()


    # DONE!
    imsave(output, img_as_ubyte(img_out))
    logging.info("Output saved to {0}.".format(output))