1 // Copyright 2019 Joe Drago. All rights reserved.
2 // SPDX-License-Identifier: BSD-2-Clause
3 
4 #include "avif/internal.h"
5 
6 #include "rav1e/rav1e.h"
7 
8 #include <string.h>
9 
10 struct avifCodecInternal
11 {
12     RaContext * rav1eContext;
13     RaChromaSampling chromaSampling;
14     int yShift;
15 };
16 
rav1eCodecDestroyInternal(avifCodec * codec)17 static void rav1eCodecDestroyInternal(avifCodec * codec)
18 {
19     if (codec->internal->rav1eContext) {
20         rav1e_context_unref(codec->internal->rav1eContext);
21         codec->internal->rav1eContext = NULL;
22     }
23     avifFree(codec->internal);
24 }
25 
26 // Official support wasn't added until v0.4.0
rav1eSupports400(void)27 static avifBool rav1eSupports400(void)
28 {
29     const char * rav1eVersionString = rav1e_version_short();
30 
31     // Check major version > 0
32     int majorVersion = atoi(rav1eVersionString);
33     if (majorVersion > 0) {
34         return AVIF_TRUE;
35     }
36 
37     // Check minor version >= 4
38     const char * minorVersionString = strchr(rav1eVersionString, '.');
39     if (!minorVersionString) {
40         return AVIF_FALSE;
41     }
42     ++minorVersionString;
43     if (!(*minorVersionString)) {
44         return AVIF_FALSE;
45     }
46     int minorVersion = atoi(minorVersionString);
47     return minorVersion >= 4;
48 }
49 
rav1eCodecEncodeImage(avifCodec * codec,avifEncoder * encoder,const avifImage * image,avifBool alpha,uint32_t addImageFlags,avifCodecEncodeOutput * output)50 static avifResult rav1eCodecEncodeImage(avifCodec * codec,
51                                         avifEncoder * encoder,
52                                         const avifImage * image,
53                                         avifBool alpha,
54                                         uint32_t addImageFlags,
55                                         avifCodecEncodeOutput * output)
56 {
57     avifResult result = AVIF_RESULT_UNKNOWN_ERROR;
58 
59     RaConfig * rav1eConfig = NULL;
60     RaFrame * rav1eFrame = NULL;
61 
62     if (!codec->internal->rav1eContext) {
63         if (codec->csOptions->count > 0) {
64             // None are currently supported!
65             return AVIF_RESULT_INVALID_CODEC_SPECIFIC_OPTION;
66         }
67 
68         const avifBool supports400 = rav1eSupports400();
69         RaPixelRange rav1eRange;
70         if (alpha) {
71             rav1eRange = (image->alphaRange == AVIF_RANGE_FULL) ? RA_PIXEL_RANGE_FULL : RA_PIXEL_RANGE_LIMITED;
72             codec->internal->chromaSampling = supports400 ? RA_CHROMA_SAMPLING_CS400 : RA_CHROMA_SAMPLING_CS420;
73             codec->internal->yShift = 1;
74         } else {
75             rav1eRange = (image->yuvRange == AVIF_RANGE_FULL) ? RA_PIXEL_RANGE_FULL : RA_PIXEL_RANGE_LIMITED;
76             codec->internal->yShift = 0;
77             switch (image->yuvFormat) {
78                 case AVIF_PIXEL_FORMAT_YUV444:
79                     codec->internal->chromaSampling = RA_CHROMA_SAMPLING_CS444;
80                     break;
81                 case AVIF_PIXEL_FORMAT_YUV422:
82                     codec->internal->chromaSampling = RA_CHROMA_SAMPLING_CS422;
83                     break;
84                 case AVIF_PIXEL_FORMAT_YUV420:
85                     codec->internal->chromaSampling = RA_CHROMA_SAMPLING_CS420;
86                     codec->internal->yShift = 1;
87                     break;
88                 case AVIF_PIXEL_FORMAT_YUV400:
89                     codec->internal->chromaSampling = supports400 ? RA_CHROMA_SAMPLING_CS400 : RA_CHROMA_SAMPLING_CS420;
90                     codec->internal->yShift = 1;
91                     break;
92                 case AVIF_PIXEL_FORMAT_NONE:
93                 default:
94                     return AVIF_RESULT_UNKNOWN_ERROR;
95             }
96         }
97 
98         rav1eConfig = rav1e_config_default();
99         if (rav1e_config_set_pixel_format(rav1eConfig,
100                                           (uint8_t)image->depth,
101                                           codec->internal->chromaSampling,
102                                           (RaChromaSamplePosition)image->yuvChromaSamplePosition,
103                                           rav1eRange) < 0) {
104             goto cleanup;
105         }
106 
107         if (addImageFlags & AVIF_ADD_IMAGE_FLAG_SINGLE) {
108             if (rav1e_config_parse(rav1eConfig, "still_picture", "true") == -1) {
109                 goto cleanup;
110             }
111         }
112         if (rav1e_config_parse_int(rav1eConfig, "width", image->width) == -1) {
113             goto cleanup;
114         }
115         if (rav1e_config_parse_int(rav1eConfig, "height", image->height) == -1) {
116             goto cleanup;
117         }
118         if (rav1e_config_parse_int(rav1eConfig, "threads", encoder->maxThreads) == -1) {
119             goto cleanup;
120         }
121 
122         int minQuantizer = AVIF_CLAMP(encoder->minQuantizer, 0, 63);
123         int maxQuantizer = AVIF_CLAMP(encoder->maxQuantizer, 0, 63);
124         if (alpha) {
125             minQuantizer = AVIF_CLAMP(encoder->minQuantizerAlpha, 0, 63);
126             maxQuantizer = AVIF_CLAMP(encoder->maxQuantizerAlpha, 0, 63);
127         }
128         minQuantizer = (minQuantizer * 255) / 63; // Rescale quantizer values as rav1e's QP range is [0,255]
129         maxQuantizer = (maxQuantizer * 255) / 63;
130         if (rav1e_config_parse_int(rav1eConfig, "min_quantizer", minQuantizer) == -1) {
131             goto cleanup;
132         }
133         if (rav1e_config_parse_int(rav1eConfig, "quantizer", maxQuantizer) == -1) {
134             goto cleanup;
135         }
136         if (encoder->tileRowsLog2 != 0) {
137             int tileRowsLog2 = AVIF_CLAMP(encoder->tileRowsLog2, 0, 6);
138             if (rav1e_config_parse_int(rav1eConfig, "tile_rows", 1 << tileRowsLog2) == -1) {
139                 goto cleanup;
140             }
141         }
142         if (encoder->tileColsLog2 != 0) {
143             int tileColsLog2 = AVIF_CLAMP(encoder->tileColsLog2, 0, 6);
144             if (rav1e_config_parse_int(rav1eConfig, "tile_cols", 1 << tileColsLog2) == -1) {
145                 goto cleanup;
146             }
147         }
148         if (encoder->speed != AVIF_SPEED_DEFAULT) {
149             int speed = AVIF_CLAMP(encoder->speed, 0, 10);
150             if (rav1e_config_parse_int(rav1eConfig, "speed", speed) == -1) {
151                 goto cleanup;
152             }
153         }
154 
155         rav1e_config_set_color_description(rav1eConfig,
156                                            (RaMatrixCoefficients)image->matrixCoefficients,
157                                            (RaColorPrimaries)image->colorPrimaries,
158                                            (RaTransferCharacteristics)image->transferCharacteristics);
159 
160         codec->internal->rav1eContext = rav1e_context_new(rav1eConfig);
161         if (!codec->internal->rav1eContext) {
162             goto cleanup;
163         }
164     }
165 
166     rav1eFrame = rav1e_frame_new(codec->internal->rav1eContext);
167 
168     int byteWidth = (image->depth > 8) ? 2 : 1;
169     if (alpha) {
170         rav1e_frame_fill_plane(rav1eFrame, 0, image->alphaPlane, image->alphaRowBytes * image->height, image->alphaRowBytes, byteWidth);
171     } else {
172         rav1e_frame_fill_plane(rav1eFrame, 0, image->yuvPlanes[0], image->yuvRowBytes[0] * image->height, image->yuvRowBytes[0], byteWidth);
173         if (image->yuvFormat != AVIF_PIXEL_FORMAT_YUV400) {
174             uint32_t uvHeight = (image->height + codec->internal->yShift) >> codec->internal->yShift;
175             rav1e_frame_fill_plane(rav1eFrame, 1, image->yuvPlanes[1], image->yuvRowBytes[1] * uvHeight, image->yuvRowBytes[1], byteWidth);
176             rav1e_frame_fill_plane(rav1eFrame, 2, image->yuvPlanes[2], image->yuvRowBytes[2] * uvHeight, image->yuvRowBytes[2], byteWidth);
177         }
178     }
179 
180     RaFrameTypeOverride frameType = RA_FRAME_TYPE_OVERRIDE_NO;
181     if (addImageFlags & AVIF_ADD_IMAGE_FLAG_FORCE_KEYFRAME) {
182         frameType = RA_FRAME_TYPE_OVERRIDE_KEY;
183     }
184     rav1e_frame_set_type(rav1eFrame, frameType);
185 
186     RaEncoderStatus encoderStatus = rav1e_send_frame(codec->internal->rav1eContext, rav1eFrame);
187     if (encoderStatus != RA_ENCODER_STATUS_SUCCESS) {
188         goto cleanup;
189     }
190 
191     RaPacket * pkt = NULL;
192     for (;;) {
193         encoderStatus = rav1e_receive_packet(codec->internal->rav1eContext, &pkt);
194         if (encoderStatus == RA_ENCODER_STATUS_ENCODED) {
195             continue;
196         }
197         if ((encoderStatus != RA_ENCODER_STATUS_SUCCESS) && (encoderStatus != RA_ENCODER_STATUS_NEED_MORE_DATA)) {
198             goto cleanup;
199         } else if (pkt) {
200             if (pkt->data && (pkt->len > 0)) {
201                 avifCodecEncodeOutputAddSample(output, pkt->data, pkt->len, (pkt->frame_type == RA_FRAME_TYPE_KEY));
202             }
203             rav1e_packet_unref(pkt);
204             pkt = NULL;
205         } else {
206             break;
207         }
208     }
209     result = AVIF_RESULT_OK;
210 cleanup:
211     if (rav1eFrame) {
212         rav1e_frame_unref(rav1eFrame);
213         rav1eFrame = NULL;
214     }
215     if (rav1eConfig) {
216         rav1e_config_unref(rav1eConfig);
217         rav1eConfig = NULL;
218     }
219     return result;
220 }
221 
rav1eCodecEncodeFinish(avifCodec * codec,avifCodecEncodeOutput * output)222 static avifBool rav1eCodecEncodeFinish(avifCodec * codec, avifCodecEncodeOutput * output)
223 {
224     for (;;) {
225         RaEncoderStatus encoderStatus = rav1e_send_frame(codec->internal->rav1eContext, NULL); // flush
226         if (encoderStatus != RA_ENCODER_STATUS_SUCCESS) {
227             return AVIF_FALSE;
228         }
229 
230         avifBool gotPacket = AVIF_FALSE;
231         RaPacket * pkt = NULL;
232         for (;;) {
233             encoderStatus = rav1e_receive_packet(codec->internal->rav1eContext, &pkt);
234             if (encoderStatus == RA_ENCODER_STATUS_ENCODED) {
235                 continue;
236             }
237             if ((encoderStatus != RA_ENCODER_STATUS_SUCCESS) && (encoderStatus != RA_ENCODER_STATUS_LIMIT_REACHED)) {
238                 return AVIF_FALSE;
239             }
240             if (pkt) {
241                 gotPacket = AVIF_TRUE;
242                 if (pkt->data && (pkt->len > 0)) {
243                     avifCodecEncodeOutputAddSample(output, pkt->data, pkt->len, (pkt->frame_type == RA_FRAME_TYPE_KEY));
244                 }
245                 rav1e_packet_unref(pkt);
246                 pkt = NULL;
247             } else {
248                 break;
249             }
250         }
251 
252         if (!gotPacket) {
253             break;
254         }
255     }
256     return AVIF_TRUE;
257 }
258 
avifCodecVersionRav1e(void)259 const char * avifCodecVersionRav1e(void)
260 {
261     return rav1e_version_full();
262 }
263 
avifCodecCreateRav1e(void)264 avifCodec * avifCodecCreateRav1e(void)
265 {
266     avifCodec * codec = (avifCodec *)avifAlloc(sizeof(avifCodec));
267     memset(codec, 0, sizeof(struct avifCodec));
268     codec->encodeImage = rav1eCodecEncodeImage;
269     codec->encodeFinish = rav1eCodecEncodeFinish;
270     codec->destroyInternal = rav1eCodecDestroyInternal;
271 
272     codec->internal = (struct avifCodecInternal *)avifAlloc(sizeof(struct avifCodecInternal));
273     memset(codec->internal, 0, sizeof(struct avifCodecInternal));
274     return codec;
275 }
276