From 473b0ba5950906c50ffd2bace023ce611aa9597d Mon Sep 17 00:00:00 2001
From: Pavel Tcherniaev <Pavel.Tcherniaev@amd.com>
Date: Thu, 16 May 2024 08:36:40 -0700
Subject: [PATCH] added MD5 checksum for videodecodemem sample (#349)

[ROCm/rocdecode commit: 746877ef563630954ed2a89abc9e5d3d57ede393]
---
 .../samples/videoDecodeMem/videodecodemem.cpp | 67 ++++++++++++++++++-
 1 file changed, 66 insertions(+), 1 deletion(-)

diff --git a/projects/rocdecode/samples/videoDecodeMem/videodecodemem.cpp b/projects/rocdecode/samples/videoDecodeMem/videodecodemem.cpp
index 160c0ae5bb..7064133cae 100644
--- a/projects/rocdecode/samples/videoDecodeMem/videodecodemem.cpp
+++ b/projects/rocdecode/samples/videoDecodeMem/videodecodemem.cpp
@@ -21,6 +21,7 @@ THE SOFTWARE.
 */
 
 #include <iostream>
+#include <fstream>
 #include <iomanip>
 #include <unistd.h>
 #include <vector>
@@ -74,6 +75,8 @@ void ShowHelpAndExit(const char *option = NULL) {
     << "-d GPU device ID (0 for the first device, 1 for the second, etc.); optional; default: 0" << 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
+    << "-md5_check MD5 File Path - generate MD5 message digest on the decoded YUV image sequence and compare to the reference MD5 string in a file; optional;" << std::endl
     << "-crop crop rectangle for output (not used when using interopped decoded frame); optional; default: 0" << std::endl
     << "-m output_surface_memory_type - decoded surface memory; optional; default - 0"
     << " [0 : OUT_SURFACE_MEM_DEV_INTERNAL/ 1 : OUT_SURFACE_MEM_DEV_COPIED/ 2 : OUT_SURFACE_MEM_HOST_COPIED/ 3 : OUT_SURFACE_MEM_NOT_MAPPED]" << std::endl;
@@ -82,11 +85,14 @@ void ShowHelpAndExit(const char *option = NULL) {
 
 int main(int argc, char **argv) {
 
-    std::string input_file_path, output_file_path;
+    std::string input_file_path, output_file_path, md5_file_path;
+    std::fstream ref_md5_file;
     int dump_output_frames = 0;
     int device_id = 0;
     bool b_force_zero_latency = false;     // false by default: enabling this option might affect decoding performance
     bool b_extract_sei_messages = false;
+    bool b_generate_md5 = false;
+    bool b_md5_check = false;
     Rect crop_rect = {};
     Rect *p_crop_rect = nullptr;
     OutputSurfaceMemoryType mem_type = OUT_SURFACE_MEM_DEV_INTERNAL;        // set to internal
@@ -134,6 +140,22 @@ int main(int argc, char **argv) {
             b_extract_sei_messages = true;
             continue;
         }
+        if (!strcmp(argv[i], "-md5")) {
+            if (i == argc) {
+                ShowHelpAndExit("-md5");
+            }
+            b_generate_md5 = true;
+            continue;
+        }
+        if (!strcmp(argv[i], "-md5_check")) {
+            if (++i == argc) {
+                ShowHelpAndExit("-md5_check");
+            }
+            b_generate_md5 = true;
+            b_md5_check = true;
+            md5_file_path = argv[i];
+            continue;
+        }
         if (!strcmp(argv[i], "-crop")) {
             if (++i == argc || 4 != sscanf(argv[i], "%d,%d,%d,%d", &crop_rect.left, &crop_rect.top, &crop_rect.right, &crop_rect.bottom)) {
                 ShowHelpAndExit("-crop");
@@ -178,6 +200,13 @@ int main(int argc, char **argv) {
         uint32_t width, height;
         double total_dec_time = 0;
 
+        if (b_generate_md5) {
+            viddec.InitMd5();
+        }
+        if (b_md5_check) {
+            ref_md5_file.open(md5_file_path.c_str(), std::ios::in);
+        }
+
         do {
             auto start_time = std::chrono::high_resolution_clock::now();
             demuxer.Demux(&pvideo, &n_video_bytes, &pts);
@@ -195,6 +224,9 @@ int main(int argc, char **argv) {
             }
             for (int i = 0; i < n_frame_returned; i++) {
                 pframe = viddec.GetFrame(&pts);
+                if (b_generate_md5) {
+                    viddec.UpdateMd5ForFrame(pframe, surf_info);
+                }
                 if (dump_output_frames && mem_type != OUT_SURFACE_MEM_NOT_MAPPED) {
                     viddec.SaveFrameToFile(output_file_path, pframe, surf_info);
                 }
@@ -216,6 +248,39 @@ int main(int argc, char **argv) {
             }
 
         }
+        if (b_generate_md5) {
+            uint8_t *digest;
+            viddec.FinalizeMd5(&digest);
+            std::cout << "MD5 message digest: ";
+            for (int i = 0; i < 16; i++) {
+                std::cout << std::setfill('0') << std::setw(2) << std::hex << static_cast<int>(digest[i]);
+            }
+            std::cout << std::endl;
+
+            if (b_md5_check) {
+                char ref_md5_string[33], c2[2];
+                uint8_t ref_md5[16];
+                std::string str;
+
+                for (int i = 0; i < 16; i++) {
+                    int c;
+                    ref_md5_file.get(c2[0]);
+                    ref_md5_file.get(c2[1]);
+                    str = c2;
+                    c = std::stoi(str, nullptr, 16);
+                    ref_md5[i] = c;
+                }
+                if (memcmp(digest, ref_md5, 16) == 0) {
+                    std::cout << "MD5 digest matches the reference MD5 digest: ";
+                } else {
+                    std::cout << "MD5 digest does not match the reference MD5 digest: ";
+                }
+                ref_md5_file.seekg(0, std::ios_base::beg);
+                ref_md5_file.getline(ref_md5_string, 33);
+                std::cout << ref_md5_string << std::endl;
+                ref_md5_file.close();
+            }
+        }
     } catch (const std::exception &ex) {
       std::cout << ex.what() << std::endl;
       exit(1);