cuvid: improve hevc profile detection

2 files changed, 30 insertions, 10 deletions

diff --git a/decoder/LAVVideo/decoders/cuvid.cpp b/decoder/LAVVideo/decoders/cuvid.cpp
index e856dac6..816319a9 100644
--- a/decoder/LAVVideo/decoders/cuvid.cpp
+++ b/decoder/LAVVideo/decoders/cuvid.cpp
@@ -625,17 +625,15 @@ STDMETHODIMP CDecCuvid::InitDecoder(AVCodecID codec, const CMediaType *pmt)
       CVC1HeaderParser vc1Parser(m_VideoParserExInfo.raw_seqhdr_data, m_VideoParserExInfo.format.seqhdr_data_length);
       m_bInterlaced = vc1Parser.hdr.interlaced;
     } else if (cudaCodec == cudaVideoCodec_HEVC) {
-      CHEVCSequenceParser hevcParser;
-      hevcParser.ParseNALs(m_VideoParserExInfo.raw_seqhdr_data, m_VideoParserExInfo.format.seqhdr_data_length, 0);
-      if (hevcParser.sps.valid) {
-        if (hevcParser.sps.profile > FF_PROFILE_HEVC_MAIN) {
-          DbgLog((LOG_TRACE, 10, L"  -> SPS indicates incompatible codec profile (profile: %d)", hevcParser.sps.profile));
-          return VFW_E_UNSUPPORTED_VIDEO;
-        }
+      hr = CheckHEVCSequence(m_VideoParserExInfo.raw_seqhdr_data, m_VideoParserExInfo.format.seqhdr_data_length);
+      if (FAILED(hr)) {
+        return VFW_E_UNSUPPORTED_VIDEO;
+      } else if (hr == S_FALSE) {
+        m_bNeedSequenceCheck = TRUE;
       }
     }
   } else {
-    m_bNeedSequenceCheck = (cudaCodec == cudaVideoCodec_H264);
+    m_bNeedSequenceCheck = (cudaCodec == cudaVideoCodec_H264 || cudaCodec == cudaVideoCodec_HEVC);
   }
 
   oVideoParserParameters.pExtVideoInfo = &m_VideoParserExInfo;
@@ -1074,6 +1072,23 @@ STDMETHODIMP CDecCuvid::CheckH264Sequence(const BYTE *buffer, int buflen)
   return S_FALSE;
 }
 
+STDMETHODIMP CDecCuvid::CheckHEVCSequence(const BYTE *buffer, int buflen)
+{
+  DbgLog((LOG_TRACE, 10, L"CDecCuvid::CheckHEVCSequence(): Checking HEVC frame for SPS"));
+  CHEVCSequenceParser hevcParser;
+  hevcParser.ParseNALs(buffer, buflen, 0);
+  if (hevcParser.sps.valid) {
+    DbgLog((LOG_TRACE, 10, L"-> SPS found"));
+    if (hevcParser.sps.profile > FF_PROFILE_HEVC_MAIN) {
+      DbgLog((LOG_TRACE, 10, L"  -> SPS indicates video incompatible with CUVID, aborting (profile: %d)", hevcParser.sps.profile));
+      return E_FAIL;
+    }
+    DbgLog((LOG_TRACE, 10, L"-> Video seems compatible with CUVID"));
+    return S_OK;
+  }
+  return S_FALSE;
+}
+
 STDMETHODIMP CDecCuvid::Decode(const BYTE *buffer, int buflen, REFERENCE_TIME rtStart, REFERENCE_TIME rtStop, BOOL bSyncPoint, BOOL bDiscontinuity)
 {
   CUresult result;
@@ -1111,8 +1126,12 @@ STDMETHODIMP CDecCuvid::Decode(const BYTE *buffer, int buflen, REFERENCE_TIME rt
     }
   }
 
-  if (m_bNeedSequenceCheck && m_VideoDecoderInfo.CodecType == cudaVideoCodec_H264) {
-    hr = CheckH264Sequence(pCuvidPacket.payload, pCuvidPacket.payload_size);
+  if (m_bNeedSequenceCheck) {
+    if (m_VideoDecoderInfo.CodecType == cudaVideoCodec_H264) {
+      hr = CheckH264Sequence(pCuvidPacket.payload, pCuvidPacket.payload_size);
+    } else if (m_VideoDecoderInfo.CodecType == cudaVideoCodec_HEVC) {
+      hr = CheckHEVCSequence(pCuvidPacket.payload, pCuvidPacket.payload_size);
+    }
     if (FAILED(hr)) {
       m_bFormatIncompatible = TRUE;
     } else if (hr == S_OK) {
diff --git a/decoder/LAVVideo/decoders/cuvid.h b/decoder/LAVVideo/decoders/cuvid.h
index d92ad8e8..7b09fa26 100644
--- a/decoder/LAVVideo/decoders/cuvid.h
+++ b/decoder/LAVVideo/decoders/cuvid.h
@@ -75,6 +75,7 @@ private:
   STDMETHODIMP FlushParser();
 
   STDMETHODIMP CheckH264Sequence(const BYTE *buffer, int buflen);
+  STDMETHODIMP CheckHEVCSequence(const BYTE *buffer, int buflen);
 
   int GetMaxGflopsGraphicsDeviceId();