4、提取H264码流中nalu

H264的NALU提取

1、nalu单元

定义nalu的存储单元，ebsp用来存储原始的包含起始码（annexb格式）的原始码流，sodb存储去除防竞争字节后的码流，prefix是3或4字节

nalu_def.h

// nalu_def.h
#pragma once

#include <cstdint>


class Nalu
{
public:
	uint8_t* ebsp;
	uint32_t ebsp_len;
	uint8_t* sodb;
	uint32_t sodb_len;
	uint8_t prefix;

	uint8_t forbidden_zero_bit;	// 1 bit
	uint8_t nal_ref_idc;		// 2 bits
	uint8_t nal_unit_type;		// 5 bits

	Nalu();

	Nalu(uint8_t* stream, uint32_t len);

	~Nalu();

	void Init(uint8_t* stream, uint32_t len);

	int ParsePrefix();

	int ParseRBSP();

	int ParseHead();

	uint8_t GetNalUnitType();

	uint8_t* GetEBSP();

	uint8_t* GetRBSP();

	uint8_t* GetSODB();

	uint32_t GetLenthEBSP();

	uint32_t GetLenthRBSP();

	uint32_t GetLenthSODB();

};

nalu_def.c

#include <string.h>
#include "nalu_def.h"

Nalu::Nalu():
	ebsp(NULL),
	ebsp_len(0),
	sodb(NULL),
	sodb_len(0),
	prefix(0),
	forbidden_zero_bit(0),
	nal_ref_idc(0),
	nal_unit_type(0)
{
}

Nalu::Nalu(uint8_t* stream, uint32_t len):Nalu()
{
	Init(stream, len);
}

Nalu::~Nalu()
{
	if (ebsp)
	{
		delete[] ebsp;
	}
	if (sodb)
	{
		delete[] sodb;
	}
}

void Nalu::Init(uint8_t* stream, uint32_t len)
{
	if (ebsp)
	{
		delete[] ebsp;
	}
	ebsp = new uint8_t[len];
	ebsp_len = len;
	memcpy(ebsp, stream, len);
}

int Nalu::ParsePrefix()
{
	prefix = 0;
	if (ebsp_len <= 3)
	{
		if (ebsp[0] == 0 && ebsp[1] == 0 && ebsp[2] == 1)
			prefix = 3;
	}
	else
	{
		if (ebsp[0] == 0 && ebsp[1] == 0 && ebsp[2] == 1)
			prefix = 3;
		else if (ebsp[0] == 0 && ebsp[1] == 0 && ebsp[2] == 0 && ebsp[3] == 1)
			prefix = 4;
	}
	return prefix == 0 ? -1 : 0;
}

int Nalu::ParseRBSP()
{
	if (ebsp_len < prefix)
		return -1;
	if(sodb)
	{
		delete[] ebsp;
	}
	sodb = new uint8_t[ebsp_len];
	sodb_len = ebsp_len;

	uint8_t* src = ebsp + prefix;
	uint8_t* end = ebsp + ebsp_len;
	uint8_t* dest_temp = sodb;
	while (src < end)
	{
		if (src < end - 3 &&
			(0x00 == src[0] && 0x00 == src[1] && 0x03 == src[2]))
		{
			*dest_temp++ = 0x00;
			*dest_temp++ = 0x00;

			src += 3;
			continue;
		}
		*dest_temp++ = *src++;
	}
	sodb_len = dest_temp - sodb;
	return 0;
}

int Nalu::ParseHead()
{
	if (ebsp_len < prefix + 1)
		return -1;
	uint8_t* head = ebsp + prefix;
	forbidden_zero_bit = (head[0] >> 7) & 0x1;
	nal_ref_idc = (head[0] >> 5) & 0x3;
	nal_unit_type = head[0] & 0x1f;
	return 0;
}

uint8_t Nalu::GetNalUnitType()
{
	return nal_unit_type;
}

uint8_t* Nalu::GetEBSP()
{
	return ebsp;
}

uint8_t* Nalu::GetRBSP()
{
	return ebsp + prefix;
}

uint8_t* Nalu::GetSODB()
{
	return sodb;
}

uint32_t Nalu::GetLenthEBSP()
{
	return ebsp_len;
}

uint32_t Nalu::GetLenthRBSP()
{
	return ebsp_len - prefix;
}

uint32_t Nalu::GetLenthSODB()
{
	return sodb_len;
}

2、nalu文件读取器

每次读取BUFF_SIZE个字节数据到缓冲区中，分割成nalu单元后存入队列，队列为空时加载数据再次存入缓冲区中进行分割

nalu_file_reader.h

#pragma once

#include <cstdio>
#include <queue>
#include "nalu_def.h"

class NaluFileReader
{
public:
	NaluFileReader();

	~NaluFileReader();

	bool OpenFile(const char* path);

	Nalu* GetNalu();

	void CloseFile();

private:

	uint8_t _status;
	FILE* _file;
	uint8_t* _buffer;
	uint32_t _size;
	std::queue<Nalu*> _nalus;
	const int32_t BUFF_SIZE = 128 * 1024;

	void ReadFromFile();

	int32_t FindStartCode(const uint8_t* buf, uint32_t size, uint8_t* prefix);

	int Splite();
};

nalu_file_reader.c

#include "nalu_file_reader.h"

#pragma warning(disable:4996)

NaluFileReader::NaluFileReader() :
	_status(0),
	_file(NULL),
	_buffer(NULL),
	_size(0)
{
	_buffer = new uint8_t[BUFF_SIZE];
}

NaluFileReader::~NaluFileReader()
{
	delete[] _buffer;
	while (!_nalus.empty())
	{
		auto nal = _nalus.front();
		_nalus.pop();
		delete nal;
	}
}

bool NaluFileReader::OpenFile(const char* path)
{
	fopen_s(&_file, path, "rb");
	_status = 0;
	return _file ? true : false;
}

Nalu* NaluFileReader::GetNalu()
{
	if (!_nalus.empty())
	{
		Nalu* nal = _nalus.front();
		_nalus.pop();
		return nal;
	}
	else if (_status)
	{
		return NULL;
	}
	else
	{
		ReadFromFile();
		Splite();
		if (!_nalus.empty())
		{
			Nalu* nal = _nalus.front();
			_nalus.pop();
			return nal;
		}
	}
	return NULL;
}

void NaluFileReader::CloseFile()
{
	fclose(_file);
	_status = 0;
}


void NaluFileReader::ReadFromFile()
{
	if (_status)
		return;
	if (_size != 0)
	{
		// need seek
		long offset = _size;
		fseek(_file, -offset, SEEK_CUR);
		_size = 0;
	}
	size_t size = fread(_buffer, 1, BUFF_SIZE, _file);
	if (size != BUFF_SIZE)
	{
		_status = 1;
	}
	_size += size;
}

int32_t NaluFileReader::FindStartCode(const uint8_t* buf, uint32_t size, uint8_t* prefix)
{
	uint32_t pos = 0;
	*prefix = 0;
	while (1)
	{
		if (pos + 4 > size) return -1;
		if (buf[pos++] != 0x0)	continue;
		if (buf[pos++] != 0x0)	continue;
		if (buf[pos] == 0x01)
		{
			*prefix = 3;
			break;
		}
		if (buf[pos++] != 0x0)	continue;
		if (buf[pos] == 0x01)
		{
			*prefix = 4;
			break;
		}
	}
	return pos + 1;
}

int NaluFileReader::Splite()
{
	uint32_t end_index = 0;
	uint8_t prefix = 0;
	uint8_t* buf_stream = _buffer;

	end_index = FindStartCode(buf_stream, _size, &prefix);
	if (prefix == 0)
	{
		return -1;
	}

	while (1)
	{
		uint8_t prefix_pre = prefix;
		end_index = FindStartCode(buf_stream + prefix, _size - prefix, &prefix);
		if (prefix == 0)
		{
			if (_status == 1)
			{
				Nalu* nalu = new Nalu(buf_stream, _size);
				_nalus.push(nalu);
				_size = 0;
			}
			break;
		}
		else
		{
			Nalu* nalu = new Nalu;
			nalu->Init(buf_stream, end_index - prefix + prefix_pre);
			_nalus.push(nalu);

			buf_stream += nalu->GetLenthEBSP();
			_size -= nalu->GetLenthEBSP();
		}
	}
	return 0;
}

3、使用读取器打印nalu信息

简单示例打印一下所有的nalu单元长度信息

#include <iostream>
#include "nalu_file_reader.h"
using namespace std;

int main()
{
    NaluFileReader nal_reader;
    nal_reader.OpenFile("F:\\test\\source.h264");
    
    while (1)
    {
        auto nal = nal_reader.GetNalu();
        if (!nal)
            break;
        nal->ParsePrefix();     // 解析前缀0x00000001
        nal->ParseRBSP();       // 去除防竞争字节
        nal->ParseHead();       // 解析nal头字段
        cout << "nal type:" << (int)nal->GetNalUnitType() << ",len:" << nal->GetLenthEBSP() << endl;
        delete nal;
    }
}