fix(hevc): Add HEVC/H.265 caption extraction support with B-frame reordering

Author: cfsmp3

src/lib_ccx/avc_functions.c

@@ -48,6 +48,7 @@ struct avc_ctx *init_avc(void)
 	ctx->cc_databufsize = 1024;
 	ctx->cc_buffer_saved = CCX_TRUE; // Was the CC buffer saved after it was last updated?
 
+	ctx->is_hevc = 0;
 	ctx->got_seq_para = 0;
 	ctx->nal_ref_idc = 0;
 	ctx->seq_parameter_set_id = 0;
@@ -87,62 +88,108 @@ struct avc_ctx *init_avc(void)
 	return ctx;
 }
 
+// HEVC NAL unit types for SEI messages
+#define HEVC_NAL_PREFIX_SEI 39
+#define HEVC_NAL_SUFFIX_SEI 40
+#define HEVC_NAL_VPS 32
+#define HEVC_NAL_SPS 33
+#define HEVC_NAL_PPS 34
+
 void do_NAL(struct encoder_ctx *enc_ctx, struct lib_cc_decode *dec_ctx, unsigned char *NAL_start, LLONG NAL_length, struct cc_subtitle *sub)
 {
 	unsigned char *NAL_stop;
-	enum ccx_avc_nal_types nal_unit_type = *NAL_start & 0x1F;
+	int nal_unit_type;
+	int nal_header_size;
+	unsigned char *payload_start;
+
+	// Determine if this is HEVC or H.264 based on NAL header
+	// H.264 NAL header: 1 byte, type in bits [4:0]
+	// HEVC NAL header: 2 bytes, type in bits [6:1] of first byte
+	if (dec_ctx->avc_ctx->is_hevc)
+	{
+		// HEVC: NAL type is in bits [6:1] of byte 0
+		nal_unit_type = (NAL_start[0] >> 1) & 0x3F;
+		nal_header_size = 2;
+	}
+	else
+	{
+		// H.264: NAL type is in bits [4:0] of byte 0
+		nal_unit_type = NAL_start[0] & 0x1F;
+		nal_header_size = 1;
+	}
 
 	NAL_stop = NAL_length + NAL_start;
-	NAL_stop = remove_03emu(NAL_start + 1, NAL_stop); // Add +1 to NAL_stop for TS, without it for MP4. Still don't know why
+	NAL_stop = remove_03emu(NAL_start + nal_header_size, NAL_stop);
+	payload_start = NAL_start + nal_header_size;
 
-	dvprint("BEGIN NAL unit type: %d length %d ref_idc: %d - Buffered captions before: %d\n",
-		nal_unit_type, NAL_stop - NAL_start - 1, dec_ctx->avc_ctx->nal_ref_idc, !dec_ctx->avc_ctx->cc_buffer_saved);
+	dvprint("BEGIN NAL unit type: %d length %d ref_idc: %d - Buffered captions before: %d (HEVC: %d)\n",
+		nal_unit_type, NAL_stop - NAL_start - nal_header_size, dec_ctx->avc_ctx->nal_ref_idc,
+		!dec_ctx->avc_ctx->cc_buffer_saved, dec_ctx->avc_ctx->is_hevc);
 
 	if (NAL_stop == NULL) // remove_03emu failed.
 	{
 		mprint("\rNotice: NAL of type %u had to be skipped because remove_03emu failed.\n", nal_unit_type);
 		return;
 	}
 
-	if (nal_unit_type == CCX_NAL_TYPE_ACCESS_UNIT_DELIMITER_9)
-	{
-		// Found Access Unit Delimiter
-	}
-	else if (nal_unit_type == CCX_NAL_TYPE_SEQUENCE_PARAMETER_SET_7)
-	{
-		// Found sequence parameter set
-		// We need this to parse NAL type 1 (CCX_NAL_TYPE_CODED_SLICE_NON_IDR_PICTURE_1)
-		dec_ctx->avc_ctx->num_nal_unit_type_7++;
-		seq_parameter_set_rbsp(dec_ctx->avc_ctx, NAL_start + 1, NAL_stop);
-		dec_ctx->avc_ctx->got_seq_para = 1;
-	}
-	else if (dec_ctx->avc_ctx->got_seq_para && (nal_unit_type == CCX_NAL_TYPE_CODED_SLICE_NON_IDR_PICTURE_1 ||
-						    nal_unit_type == CCX_NAL_TYPE_CODED_SLICE_IDR_PICTURE)) // Only if nal_unit_type=1
+	if (dec_ctx->avc_ctx->is_hevc)
 	{
-		// Found coded slice of a non-IDR picture
-		// We only need the slice header data, no need to implement
-		// slice_layer_without_partitioning_rbsp( );
-		slice_header(enc_ctx, dec_ctx, NAL_start + 1, NAL_stop, nal_unit_type, sub);
-	}
-	else if (dec_ctx->avc_ctx->got_seq_para && nal_unit_type == CCX_NAL_TYPE_SEI)
-	{
-		// Found SEI (used for subtitles)
-		// set_fts(ctx->timing); // FIXME - check this!!!
-		sei_rbsp(dec_ctx->avc_ctx, NAL_start + 1, NAL_stop);
+		// HEVC NAL unit processing
+		if (nal_unit_type == HEVC_NAL_VPS || nal_unit_type == HEVC_NAL_SPS || nal_unit_type == HEVC_NAL_PPS)
+		{
+			// Found HEVC parameter set - mark as having sequence params
+			// We don't parse HEVC SPS fully, but we need to enable SEI processing
+			dec_ctx->avc_ctx->got_seq_para = 1;
+		}
+		else if (nal_unit_type == HEVC_NAL_PREFIX_SEI || nal_unit_type == HEVC_NAL_SUFFIX_SEI)
+		{
+			// Found HEVC SEI (used for subtitles)
+			// SEI payload format is similar to H.264
+			sei_rbsp(dec_ctx->avc_ctx, payload_start, NAL_stop);
+		}
 	}
-	else if (dec_ctx->avc_ctx->got_seq_para && nal_unit_type == CCX_NAL_TYPE_PICTURE_PARAMETER_SET)
+	else
 	{
-		// Found Picture parameter set
+		// H.264 NAL unit processing (original code)
+		if (nal_unit_type == CCX_NAL_TYPE_ACCESS_UNIT_DELIMITER_9)
+		{
+			// Found Access Unit Delimiter
+		}
+		else if (nal_unit_type == CCX_NAL_TYPE_SEQUENCE_PARAMETER_SET_7)
+		{
+			// Found sequence parameter set
+			// We need this to parse NAL type 1 (CCX_NAL_TYPE_CODED_SLICE_NON_IDR_PICTURE_1)
+			dec_ctx->avc_ctx->num_nal_unit_type_7++;
+			seq_parameter_set_rbsp(dec_ctx->avc_ctx, payload_start, NAL_stop);
+			dec_ctx->avc_ctx->got_seq_para = 1;
+		}
+		else if (dec_ctx->avc_ctx->got_seq_para && (nal_unit_type == CCX_NAL_TYPE_CODED_SLICE_NON_IDR_PICTURE_1 ||
+							    nal_unit_type == CCX_NAL_TYPE_CODED_SLICE_IDR_PICTURE))
+		{
+			// Found coded slice of a non-IDR picture
+			// We only need the slice header data
+			slice_header(enc_ctx, dec_ctx, payload_start, NAL_stop, nal_unit_type, sub);
+		}
+		else if (dec_ctx->avc_ctx->got_seq_para && nal_unit_type == CCX_NAL_TYPE_SEI)
+		{
+			// Found SEI (used for subtitles)
+			sei_rbsp(dec_ctx->avc_ctx, payload_start, NAL_stop);
+		}
+		else if (dec_ctx->avc_ctx->got_seq_para && nal_unit_type == CCX_NAL_TYPE_PICTURE_PARAMETER_SET)
+		{
+			// Found Picture parameter set
+		}
 	}
+
 	if (temp_debug)
 	{
-		int len = NAL_stop - (NAL_start + 1);
+		int len = NAL_stop - payload_start;
 		dbg_print(CCX_DMT_VIDES, "\n After decoding, the actual thing was (length =%d)\n", len);
-		dump(CCX_DMT_VIDES, NAL_start + 1, len > 160 ? 160 : len, 0, 0);
+		dump(CCX_DMT_VIDES, payload_start, len > 160 ? 160 : len, 0, 0);
 	}
 
 	dvprint("END   NAL unit type: %d length %d ref_idc: %d - Buffered captions after: %d\n",
-		nal_unit_type, NAL_stop - NAL_start - 1, dec_ctx->avc_ctx->nal_ref_idc, !dec_ctx->avc_ctx->cc_buffer_saved);
+		nal_unit_type, NAL_stop - NAL_start - nal_header_size, dec_ctx->avc_ctx->nal_ref_idc, !dec_ctx->avc_ctx->cc_buffer_saved);
 }
 
 // Process inbuf bytes in buffer holding and AVC (H.264) video stream.

src/lib_ccx/avc_functions.h

@@ -9,6 +9,7 @@ struct avc_ctx
 	long cc_databufsize;
 	int cc_buffer_saved; // Was the CC buffer saved after it was last updated?
 
+	int is_hevc; // Flag to indicate HEVC (H.265) mode vs H.264
 	int got_seq_para;
 	unsigned nal_ref_idc;
 	LLONG seq_parameter_set_id;

src/lib_ccx/ccx_common_constants.h

@@ -237,7 +237,8 @@ enum ccx_bufferdata_type
 	CCX_ISDB_SUBTITLE = 8,
 	/* BUffer where cc data contain 3 byte cc_valid ccdata 1 ccdata 2 */
 	CCX_RAW_TYPE = 9,
-	CCX_DVD_SUBTITLE = 10
+	CCX_DVD_SUBTITLE = 10,
+	CCX_HEVC = 11
 };
 
 enum ccx_frame_type

src/lib_ccx/general_loop.c

@@ -779,6 +779,13 @@ int process_data(struct encoder_ctx *enc_ctx, struct lib_cc_decode *dec_ctx, str
 	else if (data_node->bufferdatatype == CCX_H264) // H.264 data from TS file
 	{
 		dec_ctx->in_bufferdatatype = CCX_H264;
+		dec_ctx->avc_ctx->is_hevc = 0;
+		got = process_avc(enc_ctx, dec_ctx, data_node->buffer, data_node->len, dec_sub);
+	}
+	else if (data_node->bufferdatatype == CCX_HEVC) // HEVC data from TS file
+	{
+		dec_ctx->in_bufferdatatype = CCX_H264; // Use same internal type for NAL processing
+		dec_ctx->avc_ctx->is_hevc = 1;
 		got = process_avc(enc_ctx, dec_ctx, data_node->buffer, data_node->len, dec_sub);
 	}
 	else if (data_node->bufferdatatype == CCX_RAW_TYPE)

src/lib_ccx/ts_functions.c

@@ -135,7 +135,7 @@ enum ccx_bufferdata_type get_buffer_type(struct cap_info *cinfo)
 	}
 	else if (cinfo->stream == CCX_STREAM_TYPE_VIDEO_HEVC)
 	{
-		return CCX_H264; // HEVC uses same buffer type as H264
+		return CCX_HEVC;
 	}
 	else if (cinfo->stream == CCX_STREAM_TYPE_PRIVATE_MPEG2 && cinfo->codec == CCX_CODEC_DVB)
 	{
@@ -390,15 +390,26 @@ void look_for_caption_data(struct ccx_demuxer *ctx, struct ts_payload *payload)
 						// Check for H.264/H.265 NAL start codes
 						if (es_data[0] == 0x00 && es_data[1] == 0x00 && es_data[2] == 0x00 && es_data[3] == 0x01)
 						{
-							// Could be H.264 or H.265 - check NAL type
-							unsigned char nal_type = es_data[4] & 0x1F;
-							if (nal_type == 7 || nal_type == 8) // SPS or PPS
+							// Check for H.264 NAL types first (1-byte header, type in bits 4:0)
+							unsigned char h264_nal_type = es_data[4] & 0x1F;
+							if (h264_nal_type == 7 || h264_nal_type == 8) // H.264 SPS or PPS
 								stream_type = CCX_STREAM_TYPE_VIDEO_H264;
+							else
+							{
+								// Check for HEVC NAL types (2-byte header, type in bits 6:1 of first byte)
+								unsigned char hevc_nal_type = (es_data[4] >> 1) & 0x3F;
+								// HEVC VPS=32, SPS=33, PPS=34, PREFIX_SEI=39, SUFFIX_SEI=40
+								// Also check for IDR (19, 20) and CRA (21) which are common first NALs
+								if (hevc_nal_type == 32 || hevc_nal_type == 33 || hevc_nal_type == 34 ||
+								    hevc_nal_type == 39 || hevc_nal_type == 40 ||
+								    hevc_nal_type == 19 || hevc_nal_type == 20 || hevc_nal_type == 21)
+									stream_type = CCX_STREAM_TYPE_VIDEO_HEVC;
+							}
 						}
 
 						mprint("PID %u detected as video stream (no PAT/PMT) - assuming %s.\n",
 						       payload->pid,
-						       stream_type == CCX_STREAM_TYPE_VIDEO_H264 ? "H.264" : "MPEG-2");
+						       stream_type == CCX_STREAM_TYPE_VIDEO_H264 ? "H.264" : (stream_type == CCX_STREAM_TYPE_VIDEO_HEVC ? "HEVC" : "MPEG-2"));
 
 						// Register this PID as a video stream that may contain captions
 						update_capinfo(ctx, payload->pid, stream_type, CCX_CODEC_ATSC_CC, 0, NULL);
@@ -661,7 +672,10 @@ int copy_payload_to_capbuf(struct cap_info *cinfo, struct ts_payload *payload)
 	int newcapbuflen;
 
 	if (cinfo->ignore == CCX_TRUE &&
-	    (cinfo->stream != CCX_STREAM_TYPE_VIDEO_MPEG2 || !ccx_options.analyze_video_stream))
+	    ((cinfo->stream != CCX_STREAM_TYPE_VIDEO_MPEG2 &&
+	      cinfo->stream != CCX_STREAM_TYPE_VIDEO_H264 &&
+	      cinfo->stream != CCX_STREAM_TYPE_VIDEO_HEVC) ||
+	     !ccx_options.analyze_video_stream))
 	{
 		return CCX_OK;
 	}
@@ -948,7 +962,10 @@ long ts_readstream(struct ccx_demuxer *ctx, struct demuxer_data **data)
 			continue;
 		}
 		else if (cinfo->ignore == CCX_TRUE &&
-			 (cinfo->stream != CCX_STREAM_TYPE_VIDEO_MPEG2 || !ccx_options.analyze_video_stream))
+			 ((cinfo->stream != CCX_STREAM_TYPE_VIDEO_MPEG2 &&
+			   cinfo->stream != CCX_STREAM_TYPE_VIDEO_H264 &&
+			   cinfo->stream != CCX_STREAM_TYPE_VIDEO_HEVC) ||
+			  !ccx_options.analyze_video_stream))
 		{
 			if (cinfo->codec_private_data)
 			{

src/lib_ccx/ts_info.c

@@ -102,7 +102,9 @@ int get_video_stream(struct ccx_demuxer *ctx)
 	struct cap_info *iter;
 	list_for_each_entry(iter, &ctx->cinfo_tree.all_stream, all_stream, struct cap_info)
 	{
-		if (iter->stream == CCX_STREAM_TYPE_VIDEO_MPEG2)
+		if (iter->stream == CCX_STREAM_TYPE_VIDEO_MPEG2 ||
+		    iter->stream == CCX_STREAM_TYPE_VIDEO_H264 ||
+		    iter->stream == CCX_STREAM_TYPE_VIDEO_HEVC)
 			return iter->pid;
 	}
 	return -1;

src/rust/src/avc/common_types.rs

@@ -9,6 +9,7 @@ pub struct AvcContextRust {
     pub cc_databufsize: usize,
     pub cc_buffer_saved: bool,
 
+    pub is_hevc: bool,
     pub got_seq_para: bool,
     pub nal_ref_idc: u32,
     pub seq_parameter_set_id: i64,
@@ -55,6 +56,7 @@ impl Default for AvcContextRust {
             cc_databufsize: 1024,
             cc_buffer_saved: true,
 
+            is_hevc: false,
             got_seq_para: false,
             nal_ref_idc: 0,
             seq_parameter_set_id: 0,