predict_phocs.py 文件源码

def main(img_dir, output_dir, pretrained_phocnet, deploy_proto, min_image_width_height, gpu_id):
    logging_format = '[%(asctime)-19s, %(name)s, %(levelname)s] %(message)s'
    logging.basicConfig(level=logging.INFO,
                        format=logging_format)
    logger = logging.getLogger('Predict PHOCs')

    if gpu_id is None:
        caffe.set_mode_cpu()
    else:
        caffe.set_mode_gpu()
        caffe.set_device(gpu_id)

    logger.info('Loading PHOCNet...')
    phocnet = caffe.Net(deploy_proto, caffe.TEST, weights=pretrained_phocnet)

    # find all images in the supplied dir
    logger.info('Found %d word images to process', len(os.listdir(img_dir)))
    word_img_list = [cv2.imread(os.path.join(img_dir, filename), cv2.CV_LOAD_IMAGE_GRAYSCALE) 
                     for filename in sorted(os.listdir(img_dir)) if filename not in ['.', '..']]
    # push images through the PHOCNet
    logger.info('Predicting PHOCs...')
    predicted_phocs = net_output_for_word_image_list(phocnet=phocnet, word_img_list=word_img_list, 
                                                    min_img_width_height=min_image_width_height)
    # save everything
    logger.info('Saving...')
    np.save(os.path.join(output_dir, 'predicted_phocs.npy'), predicted_phocs)
    logger.info('Finished')