From f8add5ad420798ff0ae023afed5bc0a8b8abed83 Mon Sep 17 00:00:00 2001
From: Kinar <kinar.ravishankar@codeandtheory.com>
Date: Mon, 13 Nov 2023 21:17:28 -0800
Subject: [PATCH] Documented the return value and added percentile to argparser

---
 .../image_classifier_benchmark.py             | 40 +++++++++----------
 1 file changed, 19 insertions(+), 21 deletions(-)

diff --git a/mediapipe/tasks/python/benchmark/vision/image_classifier/image_classifier_benchmark.py b/mediapipe/tasks/python/benchmark/vision/image_classifier/image_classifier_benchmark.py
index 502441879..d505a8fe8 100644
--- a/mediapipe/tasks/python/benchmark/vision/image_classifier/image_classifier_benchmark.py
+++ b/mediapipe/tasks/python/benchmark/vision/image_classifier/image_classifier_benchmark.py
@@ -1,17 +1,3 @@
-# Copyright 2023 The MediaPipe Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Benchmark for the image classifier task."""
 import argparse
 import time
 import numpy as np
@@ -22,16 +8,21 @@ from mediapipe.tasks.python import vision
 _IMAGE_FILE = 'burger.jpg'
 
 
-def run(model: str, n_iterations: int, delegate: python.BaseOptions.Delegate):
+def run(model: str, n_iterations: int, delegate: python.BaseOptions.Delegate,
+        percentile: float):
     """Run asynchronous inference on images and benchmark.
 
     Args:
         model: Path to the TFLite model.
         n_iterations: Number of iterations to run the benchmark.
         delegate: CPU or GPU delegate for inference.
+        percentile: Percentage for the percentiles to compute. Values must be
+            between 0 and 100 inclusive.
+    Returns:
+      The n-th percentile of the inference times.
     """
     inference_times = []
-    
+
     # Initialize the image classifier
     base_options = python.BaseOptions(model_asset_path=model, delegate=delegate)
     options = vision.ImageClassifierOptions(
@@ -48,7 +39,7 @@ def run(model: str, n_iterations: int, delegate: python.BaseOptions.Delegate):
         inference_times.append((end_time_ns - start_time_ns) / 1_000_000)
 
     classifier.close()
-    return np.percentile(inference_times, 95)
+    return np.percentile(inference_times, percentile)
 
 
 def main():
@@ -60,15 +51,22 @@ def main():
     parser.add_argument(
         '--iterations', help='Number of iterations for benchmarking.', type=int,
         default=100)
+    parser.add_argument(
+        '--percentile', help='Percentile for benchmarking statistics.',
+        type=float, default=95.0)
     args = parser.parse_args()
 
     # Run benchmark on CPU
-    cpu_time = run(args.model, args.iterations, python.BaseOptions.Delegate.CPU)
-    print(f"95th Percentile Inference Time on CPU: {cpu_time:.6f} milliseconds")
+    cpu_time = run(args.model, args.iterations, python.BaseOptions.Delegate.CPU,
+                   args.percentile)
+    print(f"{args.percentile}th Percentile Inference Time on CPU: "
+          f"{cpu_time:.6f} milliseconds")
 
     # Run benchmark on GPU
-    gpu_time = run(args.model, args.iterations, python.BaseOptions.Delegate.GPU)
-    print(f"95th Percentile Inference Time on GPU: {gpu_time:.6f} milliseconds")
+    gpu_time = run(args.model, args.iterations, python.BaseOptions.Delegate.GPU,
+                   args.percentile)
+    print(f"{args.percentile}th Percentile Inference Time on GPU: "
+          f"{gpu_time:.6f} milliseconds")
 
 
 if __name__ == '__main__':