summaryrefslogtreecommitdiff
path: root/va/va_fool_getframe.c
blob: 0b2a8d93f056d267be8fa46d4f3c7c874c967bd3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
/* The code refers to
 * http://keyj.s2000.at/files/projects/h264-src.tar.gz
 */
#include <string.h>
#include <stdio.h>

#define SLICE_NUM 4 
#define NAL_BUF_SIZE  65536  // maximum NAL unit size
#define RING_BUF_SIZE  8192  // input ring buffer size, MUST be a power of two!

typedef struct _nal_unit {
  int NumBytesInNALunit;
  int forbidden_zero_bit;
  int nal_ref_idc;
  int nal_unit_type;
  unsigned char *last_rbsp_byte;
} nal_unit;
 typedef struct _slice_header {
  int first_mb_in_slice;
} slice_header;
 
static int get_next_nal_unit(FILE *input_fp, nal_unit *nalu); 
static int get_unsigned_exp_golomb();
static void decode_slice_header(slice_header *sh);
static void input_read(FILE *input_fp, unsigned char *dest, int size);
static int input_get_bits(int bit_count);
int va_FoolGetFrame(FILE *input_fp, char *frame_buf); 

static unsigned char nal_buf[NAL_BUF_SIZE];
static unsigned char ring_buf[RING_BUF_SIZE];
static int input_remain = 0;
static int ring_pos = 0;
static int nal_pos;
static int nal_bit;
static int frame_no = 0, cur_frame_no = 0;

#define SLICE_NUM 4 
#define RING_MOD  ((RING_BUF_SIZE)-1)
#define HALF_RING ((RING_BUF_SIZE)/2)

#define gnn_advance() do { \
	ring_pos = (ring_pos+1)&RING_MOD; \
	--input_remain; \
	if (ring_pos==0) input_read(input_fp, &ring_buf[HALF_RING],HALF_RING); \
	if (ring_pos==HALF_RING) input_read(input_fp, &ring_buf[0],HALF_RING); \
} while(0)

#define gnn_add_segment(end) do { \
	int size = end-segment_start; \
	if (size>0) { \
		memcpy(&nal_buf[nalu_size],&ring_buf[segment_start],size); \
		nalu_size += size; \
	} \
	segment_start = end&RING_MOD; \
} while(0)

static int input_get_bits(int bit_count) 
{
    int res = 0;
    register unsigned int x = 
        (nal_buf[nal_pos]<<24)|
        (nal_buf[nal_pos+1]<<16)|
        (nal_buf[nal_pos+2]<<8)|
        nal_buf[nal_pos+3];

    res = (x>>(32-bit_count-nal_bit))&((1<<bit_count)-1);
    nal_bit += bit_count;
    nal_pos += nal_bit>>3;
    nal_bit &= 7;

    return res;
}

static int input_get_one_bit() 
{
    int res = (nal_buf[nal_pos]>>(7-nal_bit))&1;

    if (++nal_bit>7) {
        ++nal_pos;
        nal_bit = 0;
    }
    return res;
}

static int get_unsigned_exp_golomb() 
{
    int exp;

    for(exp = 0; !input_get_one_bit(); ++exp);
    
    if (exp)
        return (1<<exp) - 1 + input_get_bits(exp);
    else
        return 0;
}

static void decode_slice_header(slice_header *sh ) 
{
    memset((void*)sh,0,sizeof(slice_header));
    sh->first_mb_in_slice = get_unsigned_exp_golomb(); 
}

static void input_read(FILE *input_fp, unsigned char *dest, int size) 
{
    int count = fread(dest, 1, size, input_fp);

    input_remain += count;
}

static int get_next_nal_unit(FILE *input_fp, nal_unit *nalu)
{
    int i,segment_start;
    int nalu_size = 0;
    int NumBytesInRbsp = 0;

    /* search for the next NALU start
     * here is the sync that the start of the NALU is 0x00000001
     */
    for (;;) {
        if (input_remain<= 4) {
            /* clip restart */
            memset(ring_buf,0,sizeof(char)*RING_BUF_SIZE);
            memset(nal_buf,0,sizeof(char)*NAL_BUF_SIZE);

            fseek(input_fp,0,SEEK_SET);
            input_remain = 0;
            input_read(input_fp, ring_buf, RING_BUF_SIZE);
            ring_pos = 0;
            return 1;
        }
        if ((!ring_buf[ring_pos]) &&
           (!ring_buf[(ring_pos+1)&RING_MOD]) &&
           (!ring_buf[(ring_pos+2)&RING_MOD]) &&
           ( ring_buf[(ring_pos+3)&RING_MOD]==1))
            break;
        gnn_advance();
    }
    
    for(i=0;i<4;++i)
        gnn_advance();

    /* add bytes to the NALU until the end is found */
    segment_start = ring_pos;
    while (input_remain) {
        if ((!ring_buf[ring_pos]) &&
           (!ring_buf[(ring_pos+1)&RING_MOD]) &&
           (!ring_buf[(ring_pos+2)&RING_MOD]))
            break;
        ring_pos = (ring_pos+1)&RING_MOD;
        --input_remain;
        
        if (ring_pos==0) {
            gnn_add_segment(RING_BUF_SIZE);
            input_read(input_fp, &ring_buf[HALF_RING],HALF_RING);
        }

        if (ring_pos==HALF_RING) {
            gnn_add_segment(HALF_RING);
            input_read(input_fp, &ring_buf[0], HALF_RING);
        }
    }

    gnn_add_segment(ring_pos);

    /* read the NAL unit */
    nal_pos = 0; nal_bit = 0;
    nalu->forbidden_zero_bit = input_get_bits(1);
    nalu->nal_ref_idc = input_get_bits(2);
    nalu->nal_unit_type = input_get_bits(5);
    nalu->last_rbsp_byte = &nal_buf[nalu_size-1];
    nalu->NumBytesInNALunit = nalu_size; 

    return 1;
}

int va_FoolGetFrame(FILE *input_fp, char *frame_buf) 
{
    int i = 0, frame_pos = 0;
    static slice_header sh; 
    static nal_unit nalu;

    /* save the current frame number */
    cur_frame_no = frame_no;
    
    /* read the clip , here is the first frame,
     * &let the clip go on frame by frame
     */
    if (!frame_no)
        input_read(input_fp, ring_buf,RING_BUF_SIZE);

    while (get_next_nal_unit(input_fp, &nalu)) {
        if (nalu.nal_unit_type == 1 || nalu.nal_unit_type == 5) {
            decode_slice_header(&sh);
            if (0 == sh.first_mb_in_slice) {
                ++frame_no;
                frame_pos = 0;
            }
            if (frame_no > (cur_frame_no+1))
                break;
            memcpy(frame_buf+frame_pos, nal_buf+1, sizeof(char)*(nalu.NumBytesInNALunit-1));
            frame_pos += nalu.NumBytesInNALunit;
        }
    }
    
    return 1; 
}