test.py 文件源码

def main(_):

    # Definition of the paths
    weights_path = './tiny-yolo-voc.weights'
    input_img_path = './horses.jpg'
    output_image_path = './output.jpg'

    # If you do not have the checkpoint yet keep it like this! When you will run test.py for the first time it will be created automatically
    ckpt_folder_path = './ckpt/'

    # Definition of the parameters
    input_height = 416
    input_width = 416
    score_threshold = 0.3
    iou_threshold = 0.3

    # Definition of the session
    sess = tf.InteractiveSession()
    tf.global_variables_initializer().run()

    # Check for an existing checkpoint and load the weights (if it exists) or do it from binary file
    print('Looking for a checkpoint...')
    saver = tf.train.Saver()
    _ = weights_loader.load(sess,weights_path,ckpt_folder_path,saver)

    # Preprocess the input image
    print('Preprocessing...')
    preprocessed_image = preprocessing(input_img_path,input_height,input_width)

    # Compute the predictions on the input image
    print('Computing predictions...')
    predictions = inference(sess,preprocessed_image)

    # Postprocess the predictions and save the output image
    print('Postprocessing...')
    output_image = postprocessing(predictions,input_img_path,score_threshold,iou_threshold,input_height,input_width)
    cv2.imwrite(output_image_path,output_image)