1/*
2 * Copyright (C) 2014 The Android Open Source Project
3 * All rights reserved.
4 *
5 * Redistribution and use in source and binary forms, with or without
6 * modification, are permitted provided that the following conditions
7 * are met:
8 * * Redistributions of source code must retain the above copyright
9 * notice, this list of conditions and the following disclaimer.
10 * * Redistributions in binary form must reproduce the above copyright
11 * notice, this list of conditions and the following disclaimer in
12 * the documentation and/or other materials provided with the
13 * distribution.
14 *
15 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
16 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
17 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
18 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
19 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
20 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
21 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
22 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
23 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
24 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
25 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
26 * SUCH DAMAGE.
27 */
28
29#include <stdio.h>
30#include <stdint.h>
31#include <assert.h>
32#include <stdlib.h>
33
34#include "avcenc_api.h"
35#include "avcenc_int.h"
36
37// Constants.
38enum {
39    kMaxWidth         = 720,
40    kMaxHeight        = 480,
41    kMaxFrameRate     = 30,
42    kMaxBitrate       = 2048, // in kbps.
43    kInputBufferSize  = (kMaxWidth * kMaxHeight * 3) / 2, // For YUV 420 format.
44    kOutputBufferSize = kInputBufferSize,
45    kMaxDpbBuffers    = 17,
46    kIDRFrameRefreshIntervalInSec = 1,
47};
48
49
50static void *MallocCb(void * /*userData*/, int32_t size, int32_t /*attrs*/) {
51    void *ptr = calloc(size, 1);
52    return ptr;
53}
54
55static void FreeCb(void * /*userData*/, void *ptr) {
56    free(ptr);
57}
58
59static int32_t DpbAllocCb(void * /*userData*/,
60        unsigned int sizeInMbs, unsigned int numBuffers) {
61
62    size_t frameSize = (sizeInMbs << 7) * 3;
63    if(numBuffers < kMaxDpbBuffers && frameSize <= kInputBufferSize) {
64        return 1;
65    } else {
66        return 0;
67    }
68}
69
70static int32_t BindFrameCb(void *userData, int32_t index, uint8_t **yuv) {
71     assert(index < kMaxDpbBuffers);
72     uint8_t** dpbBuffer = static_cast<uint8_t**>(userData);
73     *yuv = dpbBuffer[index];
74     return 1;
75}
76
77static void UnbindFrameCb(void * /*userData*/, int32_t /*index*/) {
78}
79
80int main(int argc, char *argv[]) {
81
82    if (argc < 7) {
83        fprintf(stderr, "Usage %s <input yuv> <output file> <width> <height>"
84                        " <frame rate> <bitrate in kbps>\n", argv[0]);
85        fprintf(stderr, "Max width %d\n", kMaxWidth);
86        fprintf(stderr, "Max height %d\n", kMaxHeight);
87        fprintf(stderr, "Max framerate %d\n", kMaxFrameRate);
88        fprintf(stderr, "Max bitrate %d kbps\n", kMaxBitrate);
89        return EXIT_FAILURE;
90    }
91
92    // Read height and width.
93    int32_t width;
94    int32_t height;
95    width = atoi(argv[3]);
96    height = atoi(argv[4]);
97    if (width > kMaxWidth || height > kMaxHeight || width <= 0 || height <= 0) {
98        fprintf(stderr, "Unsupported dimensions %dx%d\n", width, height);
99        return EXIT_FAILURE;
100    }
101
102    if (width % 16 != 0 || height % 16 != 0) {
103        fprintf(stderr, "Video frame size %dx%d must be a multiple of 16\n",
104            width, height);
105        return EXIT_FAILURE;
106    }
107
108    // Read frame rate.
109    int32_t frameRate;
110    frameRate = atoi(argv[5]);
111    if (frameRate > kMaxFrameRate || frameRate <= 0) {
112        fprintf(stderr, "Unsupported frame rate %d\n", frameRate);
113        return EXIT_FAILURE;
114    }
115
116    // Read bit rate.
117    int32_t bitrate;
118    bitrate = atoi(argv[6]);
119    if (bitrate > kMaxBitrate || bitrate <= 0) {
120        fprintf(stderr, "Unsupported bitrate %d\n", bitrate);
121        return EXIT_FAILURE;
122    }
123    bitrate *= 1024; // kbps to bps.
124
125    // Open the input file.
126    FILE *fpInput = fopen(argv[1], "rb");
127    if (!fpInput) {
128        fprintf(stderr, "Could not open %s\n", argv[1]);
129        return EXIT_FAILURE;
130    }
131
132    // Open the output file.
133    FILE *fpOutput = fopen(argv[2], "wb");
134    if (!fpOutput) {
135        fprintf(stderr, "Could not open %s\n", argv[2]);
136        fclose(fpInput);
137        return EXIT_FAILURE;
138    }
139
140    // Allocate input buffer.
141    uint8_t *inputBuf = (uint8_t *)malloc(kInputBufferSize);
142    assert(inputBuf != NULL);
143
144    // Allocate output buffer.
145    uint8_t *outputBuf = (uint8_t *)malloc(kOutputBufferSize);
146    assert(outputBuf != NULL);
147
148    // Allocate dpb buffers.
149    uint8_t * dpbBuffers[kMaxDpbBuffers];
150    for (int i = 0; i < kMaxDpbBuffers; ++i) {
151        dpbBuffers[i] = (uint8_t *)malloc(kInputBufferSize);
152        assert(dpbBuffers[i] != NULL);
153    }
154
155    // Initialize the encoder parameters.
156    tagAVCEncParam encParams;
157    memset(&encParams, 0, sizeof(tagAVCEncParam));
158    encParams.rate_control = AVC_ON;
159    encParams.initQP = 0;
160    encParams.init_CBP_removal_delay = 1600;
161
162    encParams.intramb_refresh = 0;
163    encParams.auto_scd = AVC_ON;
164    encParams.out_of_band_param_set = AVC_ON;
165    encParams.poc_type = 2;
166    encParams.log2_max_poc_lsb_minus_4 = 12;
167    encParams.delta_poc_zero_flag = 0;
168    encParams.offset_poc_non_ref = 0;
169    encParams.offset_top_bottom = 0;
170    encParams.num_ref_in_cycle = 0;
171    encParams.offset_poc_ref = NULL;
172
173    encParams.num_ref_frame = 1;
174    encParams.num_slice_group = 1;
175    encParams.fmo_type = 0;
176
177    encParams.db_filter = AVC_ON;
178    encParams.disable_db_idc = 0;
179
180    encParams.alpha_offset = 0;
181    encParams.beta_offset = 0;
182    encParams.constrained_intra_pred = AVC_OFF;
183
184    encParams.data_par = AVC_OFF;
185    encParams.fullsearch = AVC_OFF;
186    encParams.search_range = 16;
187    encParams.sub_pel = AVC_OFF;
188    encParams.submb_pred = AVC_OFF;
189    encParams.rdopt_mode = AVC_OFF;
190    encParams.bidir_pred = AVC_OFF;
191
192    encParams.use_overrun_buffer = AVC_OFF;
193
194    encParams.width = width;
195    encParams.height = height;
196    encParams.bitrate = bitrate;
197    encParams.frame_rate = 1000 * frameRate;  // In frames/ms.
198    encParams.CPB_size = (uint32_t) (bitrate >> 1);
199
200    int32_t  IDRFrameRefreshIntervalInSec = kIDRFrameRefreshIntervalInSec;
201    if (IDRFrameRefreshIntervalInSec == 0) {
202        encParams.idr_period = 1;  // All I frames.
203    } else {
204        encParams.idr_period = (IDRFrameRefreshIntervalInSec * frameRate);
205    }
206
207    int32_t nMacroBlocks = ((((width + 15) >> 4) << 4) *
208            (((height + 15) >> 4) << 4)) >> 8;
209    uint32_t *sliceGroup = (uint32_t *) malloc(sizeof(uint32_t) * nMacroBlocks);
210    assert(sliceGroup != NULL);
211    for (int i = 0, idx = 0; i < nMacroBlocks; ++i) {
212        sliceGroup[i] = idx++;
213        if (idx >= encParams.num_slice_group) {
214            idx = 0;
215        }
216    }
217    encParams.slice_group = sliceGroup;
218    encParams.profile = AVC_BASELINE;
219    encParams.level = AVC_LEVEL2;
220
221    // Initialize the handle.
222    tagAVCHandle handle;
223    memset(&handle, 0, sizeof(tagAVCHandle));
224    handle.AVCObject = NULL;
225    handle.userData = dpbBuffers;
226    handle.CBAVC_DPBAlloc = DpbAllocCb;
227    handle.CBAVC_FrameBind = BindFrameCb;
228    handle.CBAVC_FrameUnbind = UnbindFrameCb;
229    handle.CBAVC_Malloc = MallocCb;
230    handle.CBAVC_Free = FreeCb;
231
232    // Initialize the encoder.
233    AVCEnc_Status status;
234    status = PVAVCEncInitialize(&handle, &encParams, NULL, NULL);
235    if (status != AVCENC_SUCCESS) {
236        fprintf(stderr, "Failed to initialize the encoder\n");
237
238        // Release resources.
239        fclose(fpInput);
240        fclose(fpOutput);
241        free(sliceGroup);
242        free(inputBuf);
243        free(outputBuf);
244        for (int i = 0; i < kMaxDpbBuffers; ++i) {
245            free(dpbBuffers[i]);
246        }
247        return EXIT_FAILURE;
248    }
249
250    // Encode Sequence Parameter Set.
251    uint32_t dataLength = kOutputBufferSize;
252    int32_t type;
253    status = PVAVCEncodeNAL(&handle, outputBuf, &dataLength, &type);
254    assert(type == AVC_NALTYPE_SPS);
255    fwrite("\x00\x00\x00\x01", 1, 4, fpOutput); // Start Code.
256    fwrite(outputBuf, 1, dataLength, fpOutput); // SPS.
257
258    // Encode Picture Paramater Set.
259    dataLength = kOutputBufferSize;
260    status = PVAVCEncodeNAL(&handle, outputBuf, &dataLength, &type);
261    assert(type == AVC_NALTYPE_PPS);
262    fwrite("\x00\x00\x00\x01", 1, 4, fpOutput); // Start Code.
263    fwrite(outputBuf, 1, dataLength, fpOutput); // PPS.
264
265    // Core loop.
266    int32_t retVal = EXIT_SUCCESS;
267    int32_t frameSize = (width * height * 3) / 2;
268    int32_t numInputFrames = 0;
269    int32_t numNalEncoded = 0;
270    bool readyForNextFrame = true;
271
272    while (1) {
273        if (readyForNextFrame == true) {
274            // Read the input frame.
275            int32_t bytesRead;
276            bytesRead = fread(inputBuf, 1, frameSize, fpInput);
277            if (bytesRead != frameSize) {
278                break; // End of file.
279            }
280
281            // Set the input frame.
282            AVCFrameIO vin;
283            memset(&vin, 0, sizeof(vin));
284            vin.height = ((height + 15) >> 4) << 4;
285            vin.pitch  = ((width  + 15) >> 4) << 4;
286            vin.coding_timestamp = (numInputFrames * 1000) / frameRate;  // in ms
287            vin.YCbCr[0] = inputBuf;
288            vin.YCbCr[1] = vin.YCbCr[0] + vin.height * vin.pitch;
289            vin.YCbCr[2] = vin.YCbCr[1] + ((vin.height * vin.pitch) >> 2);
290            vin.disp_order = numInputFrames;
291
292            status = PVAVCEncSetInput(&handle, &vin);
293            if (status == AVCENC_SUCCESS || status == AVCENC_NEW_IDR) {
294                readyForNextFrame = false;
295                ++numInputFrames;
296            } else if (status < AVCENC_SUCCESS) {
297                fprintf(stderr, "Error %d while setting input frame\n", status);
298                retVal = EXIT_FAILURE;
299                break;
300            } else {
301                fprintf(stderr, "Frame drop\n");
302                readyForNextFrame = true;
303                ++numInputFrames;
304                continue;
305            }
306        }
307
308        // Encode the input frame.
309        dataLength = kOutputBufferSize;
310        status = PVAVCEncodeNAL(&handle, outputBuf, &dataLength, &type);
311        if (status == AVCENC_SUCCESS) {
312            PVAVCEncGetOverrunBuffer(&handle);
313        } else if (status == AVCENC_PICTURE_READY) {
314            PVAVCEncGetOverrunBuffer(&handle);
315            readyForNextFrame = true;
316            AVCFrameIO recon;
317            if (PVAVCEncGetRecon(&handle, &recon) == AVCENC_SUCCESS) {
318                PVAVCEncReleaseRecon(&handle, &recon);
319            }
320        } else {
321            dataLength = 0;
322            readyForNextFrame = true;
323        }
324
325        if (status < AVCENC_SUCCESS) {
326            fprintf(stderr, "Error %d while encoding frame\n", status);
327            retVal = EXIT_FAILURE;
328            break;
329        }
330
331        numNalEncoded++;
332
333        // Write the output.
334        if (dataLength > 0) {
335            fwrite("\x00\x00\x00\x01", 1, 4, fpOutput); // Start Code.
336            fwrite(outputBuf, 1, dataLength, fpOutput); // NAL.
337            printf("NAL %d of size %d written\n", numNalEncoded, dataLength + 4);
338        }
339    }
340
341    // Close input and output file.
342    fclose(fpInput);
343    fclose(fpOutput);
344
345    // Free allocated memory.
346    free(sliceGroup);
347    free(inputBuf);
348    free(outputBuf);
349    for (int i = 0; i < kMaxDpbBuffers; ++i) {
350        free(dpbBuffers[i]);
351    }
352
353    // Close encoder instance.
354    PVAVCCleanUpEncoder(&handle);
355
356    return retVal;
357}
358