From cbc9d4213b7769eabbfd7b293dc862556d79312d Mon Sep 17 00:00:00 2001
From: Marten Richter <marten.richter@freenet.de>
Date: Sat, 9 Feb 2013 14:57:41 +0100
Subject: [PATCH] Correct h264 framecounting for field pictures

---
 demuxer.cc   | 112 +++++++++++++++++++++++++++++++++++++++++++++++++--
 demuxer.h    |  11 ++++-
 demuxerts.cc |   7 +++-
 demuxerts.h  |   2 +
 4 files changed, 126 insertions(+), 6 deletions(-)

diff --git a/demuxer.cc b/demuxer.cc
index ea1bc8f..30dae83 100644
--- a/demuxer.cc
+++ b/demuxer.cc
@@ -36,7 +36,8 @@
 #define DEMUXER_H264_ACCESS_UNIT 0x00000109
 #define DEMUXER_H264_SEQ_PARAMETER_SET 0x00000107
 #define DEMUXER_H264_SUB_ENHANCEMENT_INF 0x00000106
-
+#define DEMUXER_H264_CODED_SLICE_IDR 0x00000105
+#define DEMUXER_H264_CODED_SLICE_NON_IDR 0x00000101
 
 #define SEEK_THRESHOLD 150000 // About 1.5 seconds
 
@@ -53,6 +54,7 @@ inline    UINT getBits(UINT num_bits);
     UINT getUe();
     int getSe();
     bool isEonalu() {return eonalu;};
+    int getPos(){return pos;};
 
 protected:
     UCHAR* nalu_buf;
@@ -308,7 +310,7 @@ UINT PESPacket::findPictureHeader(bool h264) const
   }
 }
 
-UINT PESPacket::countPictureHeaders(bool h264) const
+UINT PESPacket::countPictureHeaders(bool h264, struct PictCountInfo& pinfo) const
 {
   if (size < 12) return 0;
   UINT pattern = ( ((UINT)data[ 8] << 24) |
@@ -318,12 +320,114 @@ UINT PESPacket::countPictureHeaders(bool h264) const
   UINT pos = 11;
   UINT count=0;
   if (h264) {
-	  
+	  //inspired by vdr algorithm for frame couting by Klaus Schmidinger
 	  while (pos<size)
 	  {
           pos++;
 		  pattern = (pattern << 8) | data[pos];
-          if (pattern==DEMUXER_H264_ACCESS_UNIT) count++;
+		  if ((pattern &0xFFFFFF00)==0x00000100) {
+			  UINT testpattern=(pattern& 0xFFFFFF1f);
+			  if (testpattern==DEMUXER_H264_ACCESS_UNIT) pinfo.hasaccessunit=true;
+			  else if (testpattern==DEMUXER_H264_SEQ_PARAMETER_SET ) {
+				  pinfo.hassps=true;
+				  NALUUnit nalu(data+pos-3,getSize()-pos+3);
+				  int profile=nalu.getBits(8);
+				  nalu.getBits(8); //constraints
+				  nalu.getBits(8); //level_idc
+				  nalu.getUe(); //seq_parameter_set_id
+				  int chroma=1;
+			      pinfo.separate_color_plane_flag=0;
+				  if (profile==100 || profile==110 || profile==122 || profile==144)
+				  {
+				      chroma=nalu.getUe();
+				      if (chroma==3)
+				      {
+				    	  pinfo.separate_color_plane_flag=nalu.getBits(1);
+				      }
+				      nalu.getUe(); //bit depth lume
+				      nalu.getUe(); //bit depth chrome
+				      nalu.getBits(1);
+				      if (nalu.getBits(1))
+				      {
+				    	  for (int i=0;i<8;i++){
+				    		  if (nalu.getBits(1))
+				    		  {
+				    			  if (i<6)
+				    			  {
+				    				  UINT lastscale=8;
+				    				  UINT nextscale=8;
+				    				  for (int j=0;j<16;j++) {
+				    					  if (nextscale!=0) {
+				    						  UINT delta=nalu.getSe();
+				    						  nextscale=(lastscale+delta+256)%256;
+				    					  }
+				    					  lastscale=(nextscale==0)?lastscale:nextscale;
+				    				  }
+				    			  }
+				    			  else
+				    			  {
+				    				  UINT lastscale=8;
+				    				  UINT nextscale=8;
+				    				  for (int j=0;j<64;j++) {
+				    					  if (nextscale!=0) {
+				    						  UINT delta=nalu.getSe();
+				    						  nextscale=(lastscale+delta+256)%256;
+				    					  }
+				    					  lastscale=(nextscale==0)?lastscale:nextscale;
+				    				  }
+				    			  }
+				    		  }
+				    	  }
+				      }
+				  }
+
+				  pinfo.log2_max_frame_num=nalu.getUe()+4; //log2framenum
+				  UINT temp=nalu.getUe();
+				  if (temp==0) //pict order
+					  nalu.getUe();
+				  else if (temp==1) {
+					  nalu.getBits(1);
+					  nalu.getSe();
+					  nalu.getSe();
+					  UINT temp2=nalu.getUe();
+					  for (int i=0;i<temp2;i++)
+						  nalu.getSe();
+				  }
+				  nalu.getUe(); //Num refframes
+				  nalu.getBits(1);
+				  nalu.getUe();
+				  nalu.getUe();
+				  pinfo.frame_mbs_only_flag=nalu.getBits(1);
+				  pattern=0xFFFFFF;
+				  pos+=nalu.getPos();
+
+			  } if (testpattern==DEMUXER_H264_CODED_SLICE_IDR || testpattern==DEMUXER_H264_CODED_SLICE_NON_IDR ) {
+				/*  Log::getInstance()->log("Demuxer", Log::ERR,
+				                                "Has slice %d %d %d %d %d",pinfo.hasaccessunit, pinfo.hassps,pinfo.frame_mbs_only_flag,pinfo.separate_color_plane_flag,
+				                                pinfo.log2_max_frame_num);*/
+				  if (pinfo.hasaccessunit && pinfo.hassps) {
+
+					  NALUUnit nalu(data+pos-3,getSize()-pos+3);
+					  nalu.getUe();// first_mb_in_slice
+					  nalu.getUe();//sliectype
+					  if (!pinfo.frame_mbs_only_flag) {
+						  nalu.getUe(); //pic_paramter_set_id
+						  if (pinfo.separate_color_plane_flag) nalu.getBits(2);
+						  nalu.getBits(pinfo.log2_max_frame_num);
+						  //if (!frame_mbs_only_flag)
+						  if (nalu.getBits(1)) { //field_picture
+							  if (!nalu.getBits(1)) { //bottom_field
+								  count++;
+							  }
+						  } else count++;
+
+					  } else count++;
+					  pattern=0xFFFFFF;
+					  pos+=nalu.getPos();
+					  pinfo.hasaccessunit=false;
+				  }
+			  }
+		  }
 	  }
 	  return count;
   } else {
diff --git a/demuxer.h b/demuxer.h
index 856703d..7a0c1f9 100644
--- a/demuxer.h
+++ b/demuxer.h
@@ -37,6 +37,15 @@ class DVBSubtitles;
 class Callback;
 class DrainTarget;
 
+struct PictCountInfo
+{
+	bool hassps;
+	bool hasaccessunit;
+	int log2_max_frame_num;
+	int frame_mbs_only_flag;
+	int separate_color_plane_flag;
+};
+
 class PESPacket
 {
   public:
@@ -63,7 +72,7 @@ class PESPacket
     UINT findPictureHeader(bool h264) const;
     UINT findSeqHeader(bool h264) const;
     UINT findSeqExtHeader(bool h264) const;
-    UINT countPictureHeaders(bool h264) const;
+    UINT countPictureHeaders(bool h264, struct PictCountInfo& pinfo) const;
     static const ULLONG PTS_INVALID = (1LL << 33);
 
 
diff --git a/demuxerts.cc b/demuxerts.cc
index 19fb41d..9299aa6 100644
--- a/demuxerts.cc
+++ b/demuxerts.cc
@@ -56,6 +56,9 @@ DemuxerTS::DemuxerTS(int p_vID, int p_aID, int p_subID, int p_tID)
   havechannelinfo=false;
   //doubledframerate=false;
   framereserve=0;
+
+  pinfo.hassps=false;
+  pinfo.hasaccessunit=false;
 }
 
 void DemuxerTS::flush()
@@ -85,6 +88,8 @@ tPacket.init(PESTYPE_PRIVATE_1,PESTYPE_SUBSTREAM_TELETEXTMAX);
   tActive = false;
  // doubledframerate=false;
   framereserve=0;
+  pinfo.hassps=false;
+  pinfo.hasaccessunit=false;
 }
 
 int DemuxerTS::scan(UCHAR *buf, int len)
@@ -747,7 +752,7 @@ void DemuxerTS::parseTSPacketDetails(PESPacket &packet) // Only important stuff
     {
         packetNumber++;
     }
-    UINT pictsinpacket=packet.countPictureHeaders(h264);
+    UINT pictsinpacket=packet.countPictureHeaders(h264,pinfo);
     
     UINT numpicts=0;
   /*  if (!doubledframerate)
diff --git a/demuxerts.h b/demuxerts.h
index 5e0cab3..2ce62fb 100644
--- a/demuxerts.h
+++ b/demuxerts.h
@@ -90,6 +90,8 @@ class DemuxerTS : public Demuxer
     Mutex pts_map_mutex;
     void parseTSPacketDetails(PESPacket &packet);
 
+    struct PictCountInfo pinfo;
+
    
 
 };
-- 
2.39.2