[FFmpeg-devel] [PATCH] vf_ass: add yuv overlay

[Jean First]

---
 libavfilter/vf_ass.c |   73 +++++++++++++++++++++++++++++++++++++++++++-------
 1 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/libavfilter/vf_ass.c b/libavfilter/vf_ass.c
index 7ef78f1..a3c4798 100644
--- a/libavfilter/vf_ass.c
+++ b/libavfilter/vf_ass.c
@@ -29,6 +29,7 @@
 #include <ass/ass.h>
 
 #include "libavutil/avstring.h"
+#include "libavutil/colorspace.h"
 #include "libavutil/imgutils.h"
 #include "libavutil/pixdesc.h"
 #include "drawutils.h"
@@ -39,6 +40,7 @@ typedef struct {
     ASS_Renderer *renderer;
     ASS_Track    *track;
     int hsub, vsub;
+    int is_packed_rgb;
     char *filename;
     uint8_t rgba_map[4];
     int     pix_step[4];       ///< steps per pixel for each plane of the main output
@@ -120,6 +122,11 @@ static int query_formats(AVFilterContext *ctx)
         PIX_FMT_ARGB,  PIX_FMT_RGBA,
         PIX_FMT_ABGR,  PIX_FMT_BGRA,
         PIX_FMT_RGB24, PIX_FMT_BGR24,
+        PIX_FMT_YUV444P,  PIX_FMT_YUV422P,  PIX_FMT_YUV420P,
+        PIX_FMT_YUV411P,  PIX_FMT_YUV410P,  PIX_FMT_YUV440P,
+        PIX_FMT_YUVA420P,
+        PIX_FMT_YUVJ444P, PIX_FMT_YUVJ422P, PIX_FMT_YUVJ420P,
+        PIX_FMT_YUVJ440P,
         PIX_FMT_NONE
     };
 
@@ -137,7 +144,7 @@ static int config_input(AVFilterLink *inlink)
     double dar = inlink->w / inlink->h * sar;
 
     av_image_fill_max_pixsteps(ass->pix_step, NULL, pix_desc);
-    ff_fill_rgba_map(ass->rgba_map, inlink->format);
+    ass->is_packed_rgb = ff_fill_rgba_map(ass->rgba_map, inlink->format) >= 0;
 
     ass->hsub = pix_desc->log2_chroma_w;
     ass->vsub = pix_desc->log2_chroma_h;
@@ -150,6 +157,54 @@ static int config_input(AVFilterLink *inlink)
 
 static void null_draw_slice(AVFilterLink *link, int y, int h, int slice_dir) { }
 
+/* libass stores an RGBA color in the format RRGGBBTT, where TT is the transparency level */
+#define AR(c)  ( (c)>>24)
+#define AG(c)  (((c)>>16)&0xFF)
+#define AB(c)  (((c)>>8) &0xFF)
+#define AA(c)  ((0xFF-c) &0xFF)
+
+#define Y 0
+#define U 1
+#define V 2
+#define A 3
+
+#define SET_PIXEL_YUV(picref, yuva_color, val, x, y, hsub, vsub) {           \
+    luma_pos    = ((x)          ) + ((y)          ) * picref->linesize[0]; \
+    alpha = yuva_color[A] * (val) * 129;                               \
+    picref->data[0][luma_pos]    = (alpha * yuva_color[Y] + (255*255*129 - alpha) * picref->data[0][luma_pos]   ) >> 23; \
+    if (((x) & ((1<<(hsub)) - 1)) == 0 && ((y) & ((1<<(vsub)) - 1)) == 0) {\
+        chroma_pos1 = ((x) >> (hsub)) + ((y) >> (vsub)) * picref->linesize[1]; \
+        chroma_pos2 = ((x) >> (hsub)) + ((y) >> (vsub)) * picref->linesize[2]; \
+        picref->data[1][chroma_pos1] = (alpha * yuva_color[U] + (255*255*129 - alpha) * picref->data[1][chroma_pos1]) >> 23; \
+        picref->data[2][chroma_pos2] = (alpha * yuva_color[V] + (255*255*129 - alpha) * picref->data[2][chroma_pos2]) >> 23; \
+    }\
+}
+
+static void overlay_ass_image_yuv(AVFilterBufferRef *picref, const int hsub,
+                                  const int vsub, const ASS_Image *image)
+{
+    int i, j, alpha;
+    unsigned int luma_pos, chroma_pos1, chroma_pos2;
+
+    for (; image; image = image->next) {
+        unsigned char *row = image->bitmap;
+        uint8_t yuva_color[4];
+        yuva_color[Y] = RGB_TO_Y_CCIR(AR(image->color), AG(image->color), AB(image->color));
+        yuva_color[U] = RGB_TO_U_CCIR(AR(image->color), AG(image->color), AB(image->color), 0);
+        yuva_color[V] = RGB_TO_V_CCIR(AR(image->color), AG(image->color), AB(image->color), 0);
+        yuva_color[A] = AA(image->color);
+
+        for (i = 0; i < image->h; i++) {
+            for (j = 0; j < image->w; j++) {
+                if (row[j]) {
+                    SET_PIXEL_YUV(picref, yuva_color, row[j], image->dst_x + j, image->dst_y + i, hsub, vsub);
+                }
+            }
+            row += image->stride;
+        }
+    }
+}
+
 #define R 0
 #define G 1
 #define B 2
@@ -163,14 +218,8 @@ static void null_draw_slice(AVFilterLink *link, int y, int h, int slice_dir) { }
     *(p+b_off) = (alpha * rgba_color[B] + (255*255*129 - alpha) * *(p+b_off)) >> 23;    \
 }
 
-/* libass stores an RGBA color in the format RRGGBBTT, where TT is the transparency level */
-#define AR(c)  ( (c)>>24)
-#define AG(c)  (((c)>>16)&0xFF)
-#define AB(c)  (((c)>>8) &0xFF)
-#define AA(c)  ((0xFF-c) &0xFF)
-
-static void overlay_ass_image(AVFilterBufferRef *picref, const uint8_t *rgba_map, const int pix_step,
-                              const ASS_Image *image)
+static void overlay_ass_image_rgb(AVFilterBufferRef *picref, const uint8_t *rgba_map, const int pix_step,
+                                  const ASS_Image *image)
 {
     int i, j;
     int alpha;
@@ -208,7 +257,11 @@ static void end_frame(AVFilterLink *inlink)
     if (detect_change)
         av_log(ctx, AV_LOG_DEBUG, "Change happened at time ms:%f\n", time_ms);
 
-    overlay_ass_image(picref, ass->rgba_map, ass->pix_step[0], image);
+    if (ass->is_packed_rgb) {
+        overlay_ass_image_rgb(picref, ass->rgba_map, ass->pix_step[0], image);
+    } else {
+        overlay_ass_image_yuv(picref, ass->hsub, ass->vsub, image);
+    }
 
     avfilter_draw_slice(outlink, 0, picref->video->h, 1);
     avfilter_end_frame(outlink);
-- 
1.7.9.1