From 3501eccf0ce47eaba395e3fc117d68efbefa5e74 Mon Sep 17 00:00:00 2001
From: jeffqjiangNew <142832361+jeffqjiangNew@users.noreply.github.com>
Date: Fri, 16 Feb 2024 08:44:03 -0500
Subject: [PATCH] * rocDecode: Added a command option to limit the number of
 decoded frames to a user set value to the video decoder app. (#252)

[ROCm/rocdecode commit: ad31866ee0b270aea6be4b6bb2c2a6bae3a86d74]
---
 projects/rocdecode/samples/videoDecode/README.md     |  1 +
 .../rocdecode/samples/videoDecode/videodecode.cpp    | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/projects/rocdecode/samples/videoDecode/README.md b/projects/rocdecode/samples/videoDecode/README.md
index 2808cf0470..b186baafa3 100644
--- a/projects/rocdecode/samples/videoDecode/README.md
+++ b/projects/rocdecode/samples/videoDecode/README.md
@@ -30,6 +30,7 @@ make -j
 ./videodecode -i <input video file [required]> 
               -o <output path to save decoded YUV frames [optional]> 
               -d <GPU device ID - 0:device 0 / 1:device 1/ ... [optional - default:0]>
+              -f <Number of decoded frames - specify the number of pictures to be decoded [optional]>
               -z <force_zero_latency - Decoded frames will be flushed out for display immediately [optional]>
               -sei <extract SEI messages [optional]>
               -md5 <generate MD5 message digest on the decoded YUV image sequence [optional]>
diff --git a/projects/rocdecode/samples/videoDecode/videodecode.cpp b/projects/rocdecode/samples/videoDecode/videodecode.cpp
index 6becca1d43..899d192525 100644
--- a/projects/rocdecode/samples/videoDecode/videodecode.cpp
+++ b/projects/rocdecode/samples/videoDecode/videodecode.cpp
@@ -45,6 +45,7 @@ void ShowHelpAndExit(const char *option = NULL) {
     << "-i Input File Path - required" << std::endl
     << "-o Output File Path - dumps output if requested; optional" << std::endl
     << "-d GPU device ID (0 for the first device, 1 for the second, etc.); optional; default: 0" << std::endl
+    << "-f Number of decoded frames - specify the number of pictures to be decoded; optional" << std::endl
     << "-z force_zero_latency (force_zero_latency, Decoded frames will be flushed out for display immediately); optional;" << std::endl
     << "-sei extract SEI messages; optional;" << std::endl
     << "-md5 generate MD5 message digest on the decoded YUV image sequence; optional;" << std::endl
@@ -71,6 +72,7 @@ int main(int argc, char **argv) {
     OutputSurfaceMemoryType mem_type = OUT_SURFACE_MEM_DEV_INTERNAL;      // set to internal
     ReconfigParams reconfig_params = { 0 };
     ReconfigDumpFileStruct reconfig_user_struct = { 0 };
+    uint32_t num_decoded_frames = 0;  // default value is 0, meaning decode the entire stream
 
     // Parse command-line arguments
     if(argc <= 1) {
@@ -102,6 +104,13 @@ int main(int argc, char **argv) {
             device_id = atoi(argv[i]);
             continue;
         }
+        if (!strcmp(argv[i], "-f")) {
+            if (++i == argc) {
+                ShowHelpAndExit("-d");
+            }
+            num_decoded_frames = atoi(argv[i]);
+            continue;
+        }
         if (!strcmp(argv[i], "-z")) {
             if (i == argc) {
                 ShowHelpAndExit("-z");
@@ -231,6 +240,9 @@ int main(int argc, char **argv) {
             auto time_per_decode = std::chrono::duration<double, std::milli>(end_time - start_time).count();
             total_dec_time += time_per_decode;
             n_frame += n_frame_returned;
+            if (num_decoded_frames && num_decoded_frames == n_frame) {
+                break;
+            }
         } while (n_video_bytes);
         
         n_frame += viddec.GetNumOfFlushedFrames();