mpeg2edit.c

由bmp生成mpeg2 的I_frame 数据

/*******************************************************************
                    MPEG-2 stream edit module
 *******************************************************************/
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <io.h>
#include <fcntl.h>

#include "bitstream.h"
#include "vfapi.h"

#define MPEG2EDIT_C
#include "mpeg2edit.h"

typedef struct {
	unsigned char *data;
	int length;

	void *next;
	void *prev;
} LIST_ELEMENT;

typedef struct {
	int length;
	LIST_ELEMENT *head;
} ES_SEQUENCE_HEADER;

typedef struct {
	int length;
	LIST_ELEMENT *head;
} ES_GOP_HEADER;

typedef struct {
	int length;
	LIST_ELEMENT *head;
} ES_PICTURE;

typedef struct {
	ES_PICTURE fst;
	ES_PICTURE snd;
} ES_FRAME;

typedef struct {
	LIST_ELEMENT *pack_header;
	LIST_ELEMENT *pes_packet;

	__int64 scr_length;
} PS_PACK;

typedef struct {
	PS_PACK *data;

	void *prev;
	void *next;
} PS_PACK_LIST_ELEMENT;

typedef struct {

	__int64 in;
	__int64 out;

	__int64 input_field;
	__int64 output_field;

	int stream_id;
	
	int offset;
	
	int sequence;
	int picture;

	int output;

	int level;
	int step;

	int temporal_reference;
	
	int padding;

	int rate;
	int scale;
	int fps;

	int closed_gop;
	int broken_link;
	
} PS_PROC_VIDEO_OPTION;

static const int PS_WASTE_PACK = 0;
static const int PS_OUTPUT_PACK = 1;
static const int PS_HEADER_PACK = 2;
static const int PS_HEADER_CHAIN_PACK = 4;

static const int PS_LEVEL_SEQUENCE = 0;
static const int PS_LEVEL_GOP = 1;
static const int PS_LEVEL_PICTURE = 2;
static const int PS_LEVEL_FIELD = 3;

static const int PS_STEP_START = 0;
static const int PS_STEP_SEEK = 1;
static const int PS_STEP_OUTPUT = 2;
static const int PS_STEP_END = 3;

typedef struct {
	unsigned char *data;
	int length;
} BLOCK;

typedef struct {
	__int64 head[256];
	__int64 tail[256];
	int count;
} STREAM_ID_LIST;

/* function prototypes */
int open_mpeg2edit(const char *path, MPEG2EDIT *mpeg2edit);

static void dummy_edit(void *mpeg2edit, char *out_path, __int64 in, __int64 out);
static void dummy_close(void *mpeg2edit);
static HRESULT _stdcall dummy_callback(char *out_path, DWORD percent);

static LIST_ELEMENT *new_list_element(LIST_ELEMENT *pos, int unit_size);
static void fwrite_list(LIST_ELEMENT *head, FILE *out);
static void clear_list(LIST_ELEMENT *head);
static LIST_ELEMENT *copy_list(LIST_ELEMENT *in);
static void padding_list(LIST_ELEMENT *head, int data);
static unsigned char get_byte_list(LIST_ELEMENT *head, int offset);
static void set_byte_list(LIST_ELEMENT *head, int offset, unsigned char byte);
static int find_list(LIST_ELEMENT *head, int offset, unsigned char *pattern, int length);
static int count_list_size(LIST_ELEMENT *head);

static BLOCK *new_block(unsigned char *data, int length);
static void delete_block(BLOCK *block);
static int find_block(BLOCK *current, BLOCK *next, int offset, BLOCK *pattern);
static unsigned char get_byte_block(BLOCK *current, BLOCK *next, int offset);
static void set_byte_block(BLOCK *current, BLOCK *next, int offset, unsigned char data);

typedef int (* judge_read_list_end)(int code);
static LIST_ELEMENT *read_list(BITSTREAM *in, int unit_size, judge_read_list_end func); 

static int subtract_temporal_reference(int current, int previous);

static int check_ts(BITSTREAM *in);

static int es_open(const char *path, MPEG2EDIT *mpeg2edit);
static void es_edit(void *mpeg2edit, char *out_path, __int64 in, __int64 out);
static void es_close(void *mpeg2edit);
static int es_read_sequence_header(BITSTREAM *in, ES_SEQUENCE_HEADER *out);
static int es_read_gop_header(BITSTREAM *in, ES_GOP_HEADER *out);
static int es_read_picture(BITSTREAM *in, ES_PICTURE *out);
static int es_read_frame(BITSTREAM *in, ES_FRAME *out);
static int es_judge_read_sequence_header_end(int code);
static int es_judge_read_gop_header_end(int code);
static int es_judge_read_picture_end(int code);
static void es_write_sequence_header(ES_SEQUENCE_HEADER *in, FILE *out);
static void es_write_gop_header(ES_GOP_HEADER *in, FILE *out);
static void es_write_picture(ES_PICTURE *in, FILE *out);
static void es_write_frame(ES_FRAME *in, FILE *out);
static void es_write_sequence_end_code(FILE *out);
static void es_clear_sequence_header(ES_SEQUENCE_HEADER *seq);
static void es_clear_gop_header(ES_GOP_HEADER *gop);
static void es_clear_picture(ES_PICTURE *pic);
static void es_clear_frame(ES_FRAME *frm);
static void es_padding_picture(ES_PICTURE *pic);
static void es_padding_frame(ES_FRAME *frm);
static int es_get_fps(ES_SEQUENCE_HEADER *seq);
static void es_set_timecode(ES_GOP_HEADER *gop, __int64 frame, int fps);
static int es_get_closed_gop(ES_GOP_HEADER *gop);
static int es_get_broken_link(ES_GOP_HEADER *gop);
static void es_set_broken_link(ES_GOP_HEADER *gop, int broken_link);
static void es_set_closed_gop(ES_GOP_HEADER *gop);
static int  es_get_picture_temporal_reference(ES_PICTURE *pic);
static void es_set_picture_temporal_reference(ES_PICTURE *pic, int temporal_reference);
static int es_get_picture_coding_type_picture(ES_PICTURE *pic);
static int es_get_picture_field_count(ES_PICTURE *pic);
static int es_get_temporal_reference(ES_FRAME *frm);
static void es_set_temporal_reference(ES_FRAME *frm, int temporal_reference);
static int es_get_coding_type(ES_FRAME *frm);
static int es_get_field_count(ES_FRAME *frm);

static int ps_open(const char *path, MPEG2EDIT *mpeg2edit);
static void ps_edit(void *mpeg2edit, char *out_path, __int64 in, __int64 out);
static void ps_close(void *mpeg2edit);
static void ps_write_program_end_code(FILE *out);
static int ps_read_pack(BITSTREAM *in, PS_PACK *out);
static void ps_write_pack(PS_PACK *in, FILE *out, STREAM_ID_LIST *audio_stream);
static void ps_clear_pack(PS_PACK *pack);
static void ps_copy_system_header_pack(PS_PACK *in, PS_PACK *out);
static int ps_find_system_header(PS_PACK *pack);
static int ps_get_video_stream_id(PS_PACK *pack);
static int ps_find_stream(PS_PACK *pack, int stream_id);
static __int64 ps_get_system_clock_reference(PS_PACK *pack);
static void ps_set_system_clock_reference(PS_PACK *pack, __int64 scr);
static __int64 ps_get_pack_scr_length(PS_PACK *pack);
static void ps_set_pack_scr_length(PS_PACK *pack, __int64 scr_length);
static int ps_proc_video(PS_PACK *current, PS_PACK *next, PS_PROC_VIDEO_OPTION *opt);
static int ps_judge_read_pack_header_end(int code);
static LIST_ELEMENT *ps_read_pes_packet(BITSTREAM *in);
static void ps_init_proc_video_option(PS_PROC_VIDEO_OPTION *opt, __int64 in, __int64 out);
static void ps_set_stream_id_proc_video_option(PS_PROC_VIDEO_OPTION *opt, int stream_id);
static int ps_get_current_percent(PS_PROC_VIDEO_OPTION *opt);
static PS_PACK_LIST_ELEMENT *ps_add_pack_list(PS_PACK_LIST_ELEMENT *pos, PS_PACK *data);
static void ps_write_pack_list(PS_PACK_LIST_ELEMENT *head, FILE *out, STREAM_ID_LIST *audio_list);
static void ps_clear_pack_list(PS_PACK_LIST_ELEMENT *head);
static int ps_check_pack_chain(PS_PACK_LIST_ELEMENT *tail, PS_PACK *next);
static __int64 ps_get_pack_list_scr_length(PS_PACK_LIST_ELEMENT *head);
static void ps_set_pack_list_scr(PS_PACK_LIST_ELEMENT *head, __int64 scr);
static void ps_set_clock_reference_pes_packet(LIST_ELEMENT *pes_packet, __int64 diff);
static BLOCK *ps_join_stream(PS_PACK *pack, int stream_id);
static void ps_update_stream(PS_PACK *pack, int stream_id, BLOCK *data);
static BLOCK *ps_get_pes_packet_data(LIST_ELEMENT *pes_packet);
static void ps_set_pes_packet_data(LIST_ELEMENT *pes_packet, BLOCK *data);
static int ps_get_pes_packet_data_length(LIST_ELEMENT *pes_packet);
static int ps_get_pes_packet_header_length(LIST_ELEMENT *pes_packet);
static int ps_proc_video_picture(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_proc_video_sequence(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_proc_video_end(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_proc_video_gop(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_proc_video_other(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_proc_video_extension(BLOCK *b1, BLOCK *b2, int old_offset, int new_offset, PS_PROC_VIDEO_OPTION *opt);
static int ps_check_fill_zero_video(PS_PROC_VIDEO_OPTION *opt);
static void ps_fill_zero_video(BLOCK *b1, BLOCK *b2, int from, int to);
static void ps_initialize_video_packet(PS_PACK *pack, int stream_id);
static void ps_terminate_video_packet(PS_PACK *pack, int stream_id);
static void ps_init_stream_id_list(STREAM_ID_LIST *list);
static void ps_add_stream_id_list(int id, __int64 offset, STREAM_ID_LIST *list);
static void ps_remove_stream_id_list(int id, STREAM_ID_LIST *list);
static int ps_check_stream_id_list(int id, STREAM_ID_LIST *list);
static int ps_check_empty_stream_id_list(STREAM_ID_LIST *list);
static __int64 ps_get_head_offset_stream_id_list(int id, STREAM_ID_LIST *list);
static __int64 ps_get_tail_offset_stream_id_list(int id, STREAM_ID_LIST *list);
static void ps_check_audio_stream(PS_PACK *pack, __int64 offset, STREAM_ID_LIST *audio_stream);
static void ps_trim_audio_packet(const char *path, STREAM_ID_LIST *audio_stream);
static void ps_trim_head_audio_packet(int fd, int stream_id, __int64 offset);
static void ps_trim_tail_audio_packet(int fd, int stream_id, __int64 offset);

static unsigned char *as_find_sync(unsigned char *buf, int size);
static int as_get_frame_size(unsigned char *header, int size);

/*-----------------------------------------------------------------*/
int open_mpeg2edit(const char *path, MPEG2EDIT *mpeg2edit)
{
	int n;
	int code;
	BITSTREAM *bs;

	int system;

	mpeg2edit->edit = dummy_edit;
	mpeg2edit->callback = dummy_callback;
	mpeg2edit->close = dummy_close;
	
	bs = bs_open(path);
	if(bs == NULL){
		return 0;
	}

	if(check_ts(bs)){
		return 0;
	}

	n = 0;
	system = 0;
	while(bs_next_packet_prefix(bs)){
		code = bs_read_bits(bs, 32);
		if(code >= 0x1ba){
			system += 1;
		}

		n += 1;

		if(n == 8){
			break;
		}
	}

	bs_close(bs);

	if(n < 8){
		return 0;
	}

	if(system){
		return ps_open(path, mpeg2edit);
	}
	
	return es_open(path, mpeg2edit);
}

/*-----------------------------------------------------------------*/
static void dummy_edit(void *mpeg2edit, char *out_path, __int64 in, __int64 out)
{
	return;
}

static HRESULT _stdcall dummy_callback(char *path, DWORD percent)
{
	return VF_OK; /* do nothing */
}

static void dummy_close(void *mpeg2edit)
{
	MPEG2EDIT *p;

	p = (MPEG2EDIT *)mpeg2edit;

	p->stream = NULL;
	p->edit = dummy_edit;
	p->close = dummy_close;
}

/*-----------------------------------------------------------------*/
static LIST_ELEMENT *new_list_element(LIST_ELEMENT *pos, int unit_size)
{
	LIST_ELEMENT *element;

	element = (LIST_ELEMENT *)calloc(1, sizeof(LIST_ELEMENT));
	if(element == NULL){
		return NULL;
	}
	
	element->data = (unsigned char *)malloc(unit_size);
	if(element->data == NULL){
		free(element);
		return NULL;
	}
	
	element->length = 0;

	element->prev = pos;

	if(pos != NULL){
		element->next = pos->next;
		pos->next = element;
	}

	return element;
}

static void fwrite_list(LIST_ELEMENT *head, FILE *out)
{
	if(out == NULL){
		return;
	}

	while(head){
		fwrite(head->data, 1, head->length, out);
		head = (LIST_ELEMENT *)head->next;
	}
}

static void clear_list(LIST_ELEMENT *head)
{
	LIST_ELEMENT *current;
	LIST_ELEMENT *next;

	current = head;

	while(current){
		free(current->data);
		next = (LIST_ELEMENT *)current->next;
		free(current);
		current = next;
	}
}

static LIST_ELEMENT *copy_list(LIST_ELEMENT *head)
{
	LIST_ELEMENT *r;
	LIST_ELEMENT *current;

	if(head == NULL){
		return NULL;
	}

	current = new_list_element(NULL, head->length);
	r = current;

	memcpy(current->data, head->data, head->length);
	current->length = head->length;
	
	head = (LIST_ELEMENT *)head->next;
	
	while(head){
		current = new_list_element(current, head->length);
		memcpy(current->data, head->data, head->length);
		current->length = head->length;
		head = (LIST_ELEMENT *)head->next;
	}

	return r;
}

static void padding_list(LIST_ELEMENT *head, int data)
{
	while(head){
		memset(head->data, data, head->length);
		head = (LIST_ELEMENT *)head->next;
	}
}

static unsigned char get_byte_list(LIST_ELEMENT *head, int offset)
{
	int n;

	n = 0;
	
	while(head){
		if(offset - n < head->length){
			return head->data[offset-n];
		}
		n += head->length;
		head = (LIST_ELEMENT *)head->next;
	}

	return 0;
}

static void set_byte_list(LIST_ELEMENT *head, int offset, unsigned char byte)
{
	int n;

	n = 0;

	while(head){
		if(offset - n < head->length){
			head->data[offset-n] = byte;
			return;
		}
		n += head->length;
		head = (LIST_ELEMENT *)head->next;
	}
}

static int find_list(LIST_ELEMENT *head, int offset, unsigned char *pattern, int length)
{
	int r;
	int i,n;
	
	LIST_ELEMENT *next;

	next = (LIST_ELEMENT *)head->next;

	n = 0;
	r = offset;
	while(head){
		if(r-n < head->length){
			for(i=0;i<length;i++){
				if(r-n+i < head->length){
					if(head->data[r-n+i] == pattern[i]){
						continue;
					}else{
						break;
					}
				}else if(next == NULL){
					return -1;
				}else{
					if(next->data[r-n-head->length+i] == pattern[i]){
						continue;
					}else{
						break;
					}
				}
			}
			if(i == length){
				return r;
			}
			r += 1;
		}else{
			n += head->length;
			head = next;
			if(next != NULL){
				next = (LIST_ELEMENT *)head->next;
			}
		}
	}

	return -1;
}

static int count_list_size(LIST_ELEMENT *head)
{
	int r;

	r = 0;
	
	while(head){
		r += head->length;
		head = (LIST_ELEMENT *)head->next;
	}

	return r;
}

/*-----------------------------------------------------------------*/
static LIST_ELEMENT *read_list(BITSTREAM *in, int unit_size, judge_read_list_end func)
{
	int n,code;
	
	LIST_ELEMENT *r;
	LIST_ELEMENT *current;

	current = new_list_element(NULL, unit_size);
	if(current == NULL){
		return NULL;
	}

	r = current;
	code = bs_read_bits(in, 32);

	do{
		if( (current->length+4) >= unit_size){
			current = new_list_element(current, unit_size);
			if(current == NULL){
				return 0;
			}
		}

		if( (code >> 8) == 1 ){
			n = bs_read(in, current->data+current->length, 4);
			current->length += n;
		}
		n = bs_read_next_packet_prefix(in, current->data+current->length, unit_size-current->length);
		current->length += n;
		
		code = bs_read_bits(in, 32);
		
	}while(n && (!func(code)));

	return r;
}

/*-----------------------------------------------------------------*/
static BLOCK *new_block(unsigned char *data, int length)
{
	BLOCK *r;

	r = (BLOCK *)malloc(sizeof(BLOCK));
	if(r == NULL){
		return NULL;
	}

	r->data = (unsigned char *)calloc(length, 1);
	if(r->data == NULL){
		free(r);
		return NULL;
	}

	r->length = length;
	memcpy(r->data, data, length);