Function main

samples/python/tensorflow_object_detection_api/compare_tf.py:233–244

(args)

def main(args):
    tf_infer = TensorFlowInfer(args.saved_model, args.preprocessor, args.detection_type, args.iou_threshold)
    trt_infer = TensorRTInfer(args.engine, args.preprocessor, args.detection_type, args.iou_threshold)

    trt_batcher = ImageBatcher(args.input, *trt_infer.input_spec(), max_num_images=args.num_images, preprocessor=args.preprocessor)
    tf_infer.override_input_shape(0, [1, trt_batcher.height, trt_batcher.width, 3])  # Same size input in TF as TRT
    tf_batcher = ImageBatcher(args.input, *tf_infer.input_spec(), max_num_images=args.num_images, preprocessor=args.preprocessor)

    tf_images, tf_detections = run(tf_batcher, tf_infer, "TensorFlow", args.nms_threshold)
    trt_images, trt_detections = run(trt_batcher, trt_infer, "TensorRT", args.nms_threshold)

    compare_images(tf_images, tf_detections, trt_images, trt_detections, args.output, args.annotations, args.labels, args.detection_type)


if __name__ == "__main__":