va/va_fool_getframe.c

   1 /* The code refers to
   2  * http://keyj.s2000.at/files/projects/h264-src.tar.gz
   3  */
   4 #include <string.h>
   5 #include <stdio.h>
   6
   7 #define SLICE_NUM 4
   8 #define NAL_BUF_SIZE  65536  // maximum NAL unit size
   9 #define RING_BUF_SIZE  8192  // input ring buffer size, MUST be a power of two!
  10
  11 typedef struct _nal_unit {
  12   int NumBytesInNALunit;
  13   int forbidden_zero_bit;
  14   int nal_ref_idc;
  15   int nal_unit_type;
  16   unsigned char *last_rbsp_byte;
  17 } nal_unit;
  18
  19  typedef struct _slice_header {
  20   int first_mb_in_slice;
  21 } slice_header;
  22
  23 static int get_next_nal_unit(FILE *input_fp, nal_unit *nalu);
  24 static int get_unsigned_exp_golomb();
  25 static void decode_slice_header(slice_header *sh);
  26 static void input_read(FILE *input_fp, unsigned char *dest, int size);
  27 static int input_get_bits(int bit_count);
  28
  29
  30 static unsigned char nal_buf[NAL_BUF_SIZE];
  31 static unsigned char ring_buf[RING_BUF_SIZE];
  32 static int input_remain = 0;
  33 static int ring_pos = 0;
  34 static int nal_pos;
  35 static int nal_bit;
  36 static int frame_no = 0;
  37
  38 #define RING_MOD  ((RING_BUF_SIZE)-1)
  39 #define HALF_RING ((RING_BUF_SIZE)/2)
  40
  41 #define gnn_advance() do { \
  42         ring_pos = (ring_pos+1)&RING_MOD; \
  43         --input_remain; \
  44         if (ring_pos==0) input_read(input_fp, &ring_buf[HALF_RING],HALF_RING); \
  45         if (ring_pos==HALF_RING) input_read(input_fp, &ring_buf[0],HALF_RING); \
  46 } while(0)
  47
  48 #define gnn_add_segment(end) do { \
  49         int size = end-segment_start; \
  50         if (size>0) { \
  51                 memcpy(&nal_buf[nalu_size],&ring_buf[segment_start],size); \
  52                 nalu_size += size; \
  53         } \
  54         segment_start = end&RING_MOD; \
  55 } while(0)
  56
  57 static int input_get_bits(int bit_count)
  58 {
  59     int res = 0;
  60     register unsigned int x =
  61         (nal_buf[nal_pos]<<24)|
  62         (nal_buf[nal_pos+1]<<16)|
  63         (nal_buf[nal_pos+2]<<8)|
  64         nal_buf[nal_pos+3];
  65
  66     res = (x>>(32-bit_count-nal_bit))&((1<<bit_count)-1);
  67     nal_bit += bit_count;
  68     nal_pos += nal_bit>>3;
  69     nal_bit &= 7;
  70
  71     return res;
  72 }
  73
  74 static int input_get_one_bit()
  75 {
  76     int res = (nal_buf[nal_pos]>>(7-nal_bit))&1;
  77
  78     if (++nal_bit>7) {
  79         ++nal_pos;
  80         nal_bit = 0;
  81     }
  82     return res;
  83 }
  84
  85 static int get_unsigned_exp_golomb()
  86 {
  87     int exp;
  88
  89     for(exp = 0; !input_get_one_bit(); ++exp);
  90
  91     if (exp)
  92         return (1<<exp) - 1 + input_get_bits(exp);
  93     else
  94         return 0;
  95 }
  96
  97 static void decode_slice_header(slice_header *sh )
  98 {
  99     memset((void*)sh,0,sizeof(slice_header));
 100     sh->first_mb_in_slice = get_unsigned_exp_golomb();
 101 }
 102
 103 static void input_read(FILE *input_fp, unsigned char *dest, int size)
 104 {
 105     int count = fread(dest, 1, size, input_fp);
 106
 107     input_remain += count;
 108 }
 109
 110 static int get_next_nal_unit(FILE *input_fp, nal_unit *nalu)
 111 {
 112     int i,segment_start;
 113     int nalu_size = 0;
 114     int NumBytesInRbsp = 0;
 115
 116     /* search for the next NALU start
 117      * here is the sync that the start of the NALU is 0x00000001
 118      */
 119     for (;;) {
 120         if (input_remain<= 4) {
 121             /* clip restart */
 122             memset(ring_buf,0,sizeof(char)*RING_BUF_SIZE);
 123             memset(nal_buf,0,sizeof(char)*NAL_BUF_SIZE);
 124
 125             fseek(input_fp,0,SEEK_SET);
 126             input_remain = 0;
 127             input_read(input_fp, ring_buf, RING_BUF_SIZE);
 128             ring_pos = 0;
 129             return 1;
 130         }
 131         if ((!ring_buf[ring_pos]) &&
 132            (!ring_buf[(ring_pos+1)&RING_MOD]) &&
 133            (!ring_buf[(ring_pos+2)&RING_MOD]) &&
 134            ( ring_buf[(ring_pos+3)&RING_MOD]==1))
 135             break;
 136         gnn_advance();
 137     }
 138
 139     for(i=0;i<4;++i)
 140         gnn_advance();
 141
 142     /* add bytes to the NALU until the end is found */
 143     segment_start = ring_pos;
 144     while (input_remain) {
 145         if ((!ring_buf[ring_pos]) &&
 146            (!ring_buf[(ring_pos+1)&RING_MOD]) &&
 147            (!ring_buf[(ring_pos+2)&RING_MOD]))
 148             break;
 149         ring_pos = (ring_pos+1)&RING_MOD;
 150         --input_remain;
 151
 152         if (ring_pos==0) {
 153             gnn_add_segment(RING_BUF_SIZE);
 154             input_read(input_fp, &ring_buf[HALF_RING],HALF_RING);
 155         }
 156
 157         if (ring_pos==HALF_RING) {
 158             gnn_add_segment(HALF_RING);
 159             input_read(input_fp, &ring_buf[0], HALF_RING);
 160         }
 161     }
 162
 163     gnn_add_segment(ring_pos);
 164
 165     /* read the NAL unit */
 166     nal_pos = 0; nal_bit = 0;
 167     nalu->forbidden_zero_bit = input_get_bits(1);
 168     nalu->nal_ref_idc = input_get_bits(2);
 169     nalu->nal_unit_type = input_get_bits(5);
 170     nalu->last_rbsp_byte = &nal_buf[nalu_size-1];
 171     nalu->NumBytesInNALunit = nalu_size;
 172
 173     return 1;
 174 }
 175
 176 int va_FoolGetFrame(FILE *input_fp, char *frame_buf)
 177 {
 178     int i = 0, frame_pos = 0;
 179     static slice_header sh;
 180     static nal_unit nalu;
 181     char nal_head[4] = {0x00,0x00,0x00,0x01};
 182
 183     /* read the clip , here is the first frame,
 184      * &let the clip go on frame by frame
 185      */
 186     if (!frame_no)
 187         input_read(input_fp, ring_buf,RING_BUF_SIZE);
 188
 189     while (get_next_nal_unit(input_fp, &nalu)) {
 190         if (nalu.nal_unit_type == 7 || nalu.nal_unit_type == 8) {
 191             memcpy(frame_buf+frame_pos, nal_head, sizeof(char)*4);
 192                     frame_pos = frame_pos + 4;
 193             memcpy(frame_buf+frame_pos, nal_buf, sizeof(char)*(nalu.NumBytesInNALunit));
 194             frame_pos += nalu.NumBytesInNALunit;
 195         }
 196         else if (nalu.nal_unit_type == 1 || nalu.nal_unit_type == 5) {
 197             decode_slice_header(&sh);
 198             if (0 == sh.first_mb_in_slice) {
 199                 ++frame_no;
 200             }
 201             memcpy(frame_buf+frame_pos, nal_head, sizeof(char)*4);
 202                     frame_pos = frame_pos + 4;
 203             memcpy(frame_buf+frame_pos, nal_buf, sizeof(char)*(nalu.NumBytesInNALunit));
 204             frame_pos += nalu.NumBytesInNALunit;
 205                     break;
 206         }
 207     }
 208
 209     return frame_pos;
 210 }