predict_video.py 文件源码

def process_frame(frame_idx, img, model, write_to_dir, conf_threshold, input_size=224):
    """Finds bounding boxes in a video frame, draws these bounding boxes
    and saves the result to HDD.
    """
    # find BBs in frame
    bbs, time_model = find_bbs(img, model, conf_threshold, input_size=input_size)

    # draw BBs
    img_out = np.copy(img)
    for (bb, score) in bbs:
        if score > conf_threshold and bb.width > 2 and bb.height > 2:
            img_out = bb.draw_on_image(img_out, color=[0, 255, 0], thickness=3)

    # save to output directory
    save_to_fp = os.path.join(write_to_dir, "%05d.jpg" % (frame_idx,))
    misc.imsave(save_to_fp, img_out)

    return time_model