From 5f779a6eb712fe4edd1618e8406ba772f590ca44 Mon Sep 17 00:00:00 2001
From: jeffqjiangNew <142832361+jeffqjiangNew@users.noreply.github.com>
Date: Mon, 3 Mar 2025 09:31:49 -0500
Subject: [PATCH] * rocDecoder/Bitstream reader: Added stream type detection
 optimization. (#522)

- During stream type evaluation, when we have a high confidence score, we stop and finalize, to avoid unnecessary exhaustive type search.
---
 CHANGELOG.md                        | 2 ++
 src/bit_stream_reader/es_reader.cpp | 7 +++++--
 src/bit_stream_reader/es_reader.h   | 1 +
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8d98ffdfd4..38796bbb71 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ Full documentation for rocDecode is available at [https://rocm.docs.amd.com/proj
 
 ### Changed
 
+* Bitstream type detection optimization in bitstream reader.
+
 ### Removed
 
 ## rocDecode 0.10.0 for ROCm 6.4
diff --git a/src/bit_stream_reader/es_reader.cpp b/src/bit_stream_reader/es_reader.cpp
index 33f689d288..e22820c2b7 100644
--- a/src/bit_stream_reader/es_reader.cpp
+++ b/src/bit_stream_reader/es_reader.cpp
@@ -585,6 +585,9 @@ int RocVideoESParser::ProbeStreamType() {
                 break;
             }
         }
+        if (stream_type_score >= STREAM_TYPE_HIGH_CONFIDENCE_SCORE_THRESHOLD) {
+            break;
+        }
     }
 
     if (stream_buf) {
@@ -688,7 +691,7 @@ int RocVideoESParser::CheckAvcEStream(uint8_t *p_stream, int stream_size) {
     if (num_start_codes == 0) {
         score = 0;
     } else {
-        score = sps_present * 25 + pps_present * 25 + idr_slice_present * 15 + slice_present * 15 + first_slice_present * 15;
+        score = sps_present * 25 + pps_present * 25 + idr_slice_present * 20 + slice_present * 15 + first_slice_present * 15;
     }
     return score;
 }
@@ -832,7 +835,7 @@ int RocVideoESParser::CheckHevcEStream(uint8_t *p_stream, int stream_size) {
     if (num_start_codes == 0) {
         score = 0;
     } else {
-        score = vps_present * 20 + sps_present * 20 + pps_present * 20 + rap_slice_present * 15 + slice_present * 15 + first_slice_present * 15;
+        score = vps_present * 15 + sps_present * 20 + pps_present * 20 + rap_slice_present * 15 + slice_present * 15 + first_slice_present * 15;
     }
     return score;
 }
diff --git a/src/bit_stream_reader/es_reader.h b/src/bit_stream_reader/es_reader.h
index 11c46bddea..1e2f0e7ef4 100644
--- a/src/bit_stream_reader/es_reader.h
+++ b/src/bit_stream_reader/es_reader.h
@@ -42,6 +42,7 @@ enum {
 
 #define STREAM_PROBE_SIZE 2 * 1024
 #define STREAM_TYPE_SCORE_THRESHOLD 50
+#define STREAM_TYPE_HIGH_CONFIDENCE_SCORE_THRESHOLD 90
 
 class RocVideoESParser {
     public: