| @@ -0,0 +1,380 @@ |
| 1 | +/* |
| 2 | + * Copyright (c) 2021 Paul B Mahol |
| 3 | + * |
| 4 | + * This file is part of FFmpeg. |
| 5 | + * |
| 6 | + * FFmpeg is free software; you can redistribute it and/or |
| 7 | + * modify it under the terms of the GNU Lesser General Public |
| 8 | + * License as published by the Free Software Foundation; either |
| 9 | + * version 2.1 of the License, or (at your option) any later version. |
| 10 | + * |
| 11 | + * FFmpeg is distributed in the hope that it will be useful, |
| 12 | + * but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 13 | + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
| 14 | + * Lesser General Public License for more details. |
| 15 | + * |
| 16 | + * You should have received a copy of the GNU Lesser General Public |
| 17 | + * License along with FFmpeg; if not, write to the Free Software |
| 18 | + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
| 19 | + */ |
| 20 | + |
| 21 | +#include <float.h> |
| 22 | + |
| 23 | +#include "libavutil/opt.h" |
| 24 | +#include "libavutil/imgutils.h" |
| 25 | +#include "libavutil/intreadwrite.h" |
| 26 | +#include "avfilter.h" |
| 27 | +#include "formats.h" |
| 28 | +#include "internal.h" |
| 29 | +#include "video.h" |
| 30 | + |
| 31 | +typedef struct HSVKeyContext { |
| 32 | + const AVClass *class; |
| 33 | + |
| 34 | + float hue, hue_opt, sat, val; |
| 35 | + float similarity; |
| 36 | + float blend; |
| 37 | + |
| 38 | + float scale; |
| 39 | + |
| 40 | + float half; |
| 41 | + |
| 42 | + int depth; |
| 43 | + int max; |
| 44 | + |
| 45 | + int hsub_log2; |
| 46 | + int vsub_log2; |
| 47 | + |
| 48 | + int (*do_slice)(AVFilterContext *ctx, void *arg, |
| 49 | + int jobnr, int nb_jobs); |
| 50 | +} HSVKeyContext; |
| 51 | + |
| 52 | +#define SQR(x) ((x)*(x)) |
| 53 | + |
| 54 | +static int do_hsvkey_pixel(HSVKeyContext *s, int y, int u, int v, |
| 55 | + float hue_key, float sat_key, float val_key) |
| 56 | +{ |
| 57 | + const float similarity = s->similarity; |
| 58 | + const float scale = s->scale; |
| 59 | + const float blend = s->blend; |
| 60 | + const int imax = s->max; |
| 61 | + const float max = imax; |
| 62 | + const float half = s->half; |
| 63 | + const float uf = u - half; |
| 64 | + const float vf = v - half; |
| 65 | + const float hue = hue_key < 0.f ? -hue_key : atan2f(uf, vf) + M_PI; |
| 66 | + const float sat = sat_key < 0.f ? -sat_key : sqrtf((uf * uf + vf * vf) / (half * half * 2.f)); |
| 67 | + const float val = val_key < 0.f ? -val_key : scale * y; |
| 68 | + float diff; |
| 69 | + |
| 70 | + hue_key = fabsf(hue_key); |
| 71 | + sat_key = fabsf(sat_key); |
| 72 | + val_key = fabsf(val_key); |
| 73 | + |
| 74 | + diff = sqrtf(fmaxf(SQR(sat) * SQR(val) + |
| 75 | + SQR(sat_key) * SQR(val_key) - |
| 76 | + 2.f * sat * val * sat_key * val_key * cosf(hue_key - hue) + |
| 77 | + SQR(val - val_key), 0.f)); |
| 78 | + if (diff < similarity) { |
| 79 | + return 0; |
| 80 | + } else if (blend > FLT_MIN) { |
| 81 | + return av_clipf((diff - similarity) / blend, 0.f, 1.f) * max; |
| 82 | + } else { |
| 83 | + return imax; |
| 84 | + } |
| 85 | + |
| 86 | + return 0; |
| 87 | +} |
| 88 | + |
| 89 | +static int do_hsvkey_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs) |
| 90 | +{ |
| 91 | + HSVKeyContext *s = avctx->priv; |
| 92 | + AVFrame *frame = arg; |
| 93 | + const int slice_start = (frame->height * jobnr) / nb_jobs; |
| 94 | + const int slice_end = (frame->height * (jobnr + 1)) / nb_jobs; |
| 95 | + const int hsub_log2 = s->hsub_log2; |
| 96 | + const int vsub_log2 = s->vsub_log2; |
| 97 | + const float hue = s->hue; |
| 98 | + const float sat = s->sat; |
| 99 | + const float val = s->val; |
| 100 | + |
| 101 | + for (int y = slice_start; y < slice_end; y++) { |
| 102 | + for (int x = 0; x < frame->width; x++) { |
| 103 | + int Y = frame->data[0][frame->linesize[0] * y + x]; |
| 104 | + int u = frame->data[1][frame->linesize[1] * (y >> vsub_log2) + (x >> hsub_log2)]; |
| 105 | + int v = frame->data[2][frame->linesize[2] * (y >> vsub_log2) + (x >> hsub_log2)]; |
| 106 | + |
| 107 | + frame->data[3][frame->linesize[3] * y + x] = do_hsvkey_pixel(s, Y, u, v, hue, sat, val); |
| 108 | + } |
| 109 | + } |
| 110 | + |
| 111 | + return 0; |
| 112 | +} |
| 113 | + |
| 114 | +static int do_hsvkey16_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs) |
| 115 | +{ |
| 116 | + HSVKeyContext *s = avctx->priv; |
| 117 | + AVFrame *frame = arg; |
| 118 | + const int slice_start = (frame->height * jobnr) / nb_jobs; |
| 119 | + const int slice_end = (frame->height * (jobnr + 1)) / nb_jobs; |
| 120 | + const int hsub_log2 = s->hsub_log2; |
| 121 | + const int vsub_log2 = s->vsub_log2; |
| 122 | + const float hue = s->hue; |
| 123 | + const float sat = s->sat; |
| 124 | + const float val = s->val; |
| 125 | + |
| 126 | + for (int y = slice_start; y < slice_end; ++y) { |
| 127 | + for (int x = 0; x < frame->width; ++x) { |
| 128 | + uint16_t *dst = (uint16_t *)(frame->data[3] + frame->linesize[3] * y); |
| 129 | + int Y = AV_RN16(&frame->data[0][frame->linesize[0] * y + 2 * x]); |
| 130 | + int u = AV_RN16(&frame->data[1][frame->linesize[1] * (y >> vsub_log2) + 2 * (x >> hsub_log2)]); |
| 131 | + int v = AV_RN16(&frame->data[2][frame->linesize[2] * (y >> vsub_log2) + 2 * (x >> hsub_log2)]); |
| 132 | + |
| 133 | + dst[x] = do_hsvkey_pixel(s, Y, u, v, hue, sat, val); |
| 134 | + } |
| 135 | + } |
| 136 | + |
| 137 | + return 0; |
| 138 | +} |
| 139 | + |
| 140 | +static int do_hsvhold_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs) |
| 141 | +{ |
| 142 | + HSVKeyContext *s = avctx->priv; |
| 143 | + AVFrame *frame = arg; |
| 144 | + const int hsub_log2 = s->hsub_log2; |
| 145 | + const int vsub_log2 = s->vsub_log2; |
| 146 | + const int width = frame->width >> hsub_log2; |
| 147 | + const int height = frame->height >> vsub_log2; |
| 148 | + const int slice_start = (height * jobnr) / nb_jobs; |
| 149 | + const int slice_end = (height * (jobnr + 1)) / nb_jobs; |
| 150 | + const float scale = s->scale; |
| 151 | + const float hue = s->hue; |
| 152 | + const float sat = s->sat; |
| 153 | + const float val = s->val; |
| 154 | + |
| 155 | + for (int y = slice_start; y < slice_end; ++y) { |
| 156 | + for (int x = 0; x < width; ++x) { |
| 157 | + uint8_t *dstu = frame->data[1] + frame->linesize[1] * y; |
| 158 | + uint8_t *dstv = frame->data[2] + frame->linesize[2] * y; |
| 159 | + int Y = frame->data[0][frame->linesize[0] * (y << vsub_log2) + (x << hsub_log2)]; |
| 160 | + int u = frame->data[1][frame->linesize[1] * y + x]; |
| 161 | + int v = frame->data[2][frame->linesize[2] * y + x]; |
| 162 | + int t = do_hsvkey_pixel(s, Y, u, v, hue, sat, val); |
| 163 | + |
| 164 | + if (t > 0) { |
| 165 | + float f = 1.f - t * scale; |
| 166 | + |
| 167 | + dstu[x] = 128 + (u - 128) * f; |
| 168 | + dstv[x] = 128 + (v - 128) * f; |
| 169 | + } |
| 170 | + } |
| 171 | + } |
| 172 | + |
| 173 | + return 0; |
| 174 | +} |
| 175 | + |
| 176 | +static int do_hsvhold16_slice(AVFilterContext *avctx, void *arg, int jobnr, int nb_jobs) |
| 177 | +{ |
| 178 | + HSVKeyContext *s = avctx->priv; |
| 179 | + AVFrame *frame = arg; |
| 180 | + const int hsub_log2 = s->hsub_log2; |
| 181 | + const int vsub_log2 = s->vsub_log2; |
| 182 | + const int width = frame->width >> hsub_log2; |
| 183 | + const int height = frame->height >> vsub_log2; |
| 184 | + const int slice_start = (height * jobnr) / nb_jobs; |
| 185 | + const int slice_end = (height * (jobnr + 1)) / nb_jobs; |
| 186 | + const float scale = s->scale; |
| 187 | + const float half = s->half; |
| 188 | + const float hue = s->hue; |
| 189 | + const float sat = s->sat; |
| 190 | + const float val = s->val; |
| 191 | + |
| 192 | + for (int y = slice_start; y < slice_end; ++y) { |
| 193 | + for (int x = 0; x < width; ++x) { |
| 194 | + uint16_t *dstu = (uint16_t *)(frame->data[1] + frame->linesize[1] * y); |
| 195 | + uint16_t *dstv = (uint16_t *)(frame->data[2] + frame->linesize[2] * y); |
| 196 | + int Y = AV_RN16(&frame->data[0][frame->linesize[0] * (y << vsub_log2) + 2 * (x << hsub_log2)]); |
| 197 | + int u = AV_RN16(&frame->data[1][frame->linesize[1] * y + 2 * x]); |
| 198 | + int v = AV_RN16(&frame->data[2][frame->linesize[2] * y + 2 * x]); |
| 199 | + int t = do_hsvkey_pixel(s, Y, u, v, hue, sat, val); |
| 200 | + |
| 201 | + if (t > 0) { |
| 202 | + float f = 1.f - t * scale; |
| 203 | + |
| 204 | + dstu[x] = half + (u - half) * f; |
| 205 | + dstv[x] = half + (v - half) * f; |
| 206 | + } |
| 207 | + } |
| 208 | + } |
| 209 | + |
| 210 | + return 0; |
| 211 | +} |
| 212 | + |
| 213 | +static int filter_frame(AVFilterLink *link, AVFrame *frame) |
| 214 | +{ |
| 215 | + AVFilterContext *avctx = link->dst; |
| 216 | + HSVKeyContext *s = avctx->priv; |
| 217 | + int res; |
| 218 | + |
| 219 | + s->hue = FFSIGN(s->hue_opt) *M_PI * fmodf(526.f - fabsf(s->hue_opt), 360.f) / 180.f; |
| 220 | + if (res = ff_filter_execute(avctx, s->do_slice, frame, NULL, |
| 221 | + FFMIN(frame->height, ff_filter_get_nb_threads(avctx)))) |
| 222 | + return res; |
| 223 | + |
| 224 | + return ff_filter_frame(avctx->outputs[0], frame); |
| 225 | +} |
| 226 | + |
| 227 | +static av_cold int config_output(AVFilterLink *outlink) |
| 228 | +{ |
| 229 | + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(outlink->format); |
| 230 | + AVFilterContext *avctx = outlink->src; |
| 231 | + HSVKeyContext *s = avctx->priv; |
| 232 | + |
| 233 | + s->depth = desc->comp[0].depth; |
| 234 | + s->max = (1 << s->depth) - 1; |
| 235 | + s->half = 0.5f * s->max; |
| 236 | + s->scale = 1.f / s->max; |
| 237 | + |
| 238 | + if (!strcmp(avctx->filter->name, "hsvkey")) { |
| 239 | + s->do_slice = s->depth <= 8 ? do_hsvkey_slice : do_hsvkey16_slice; |
| 240 | + } else { |
| 241 | + s->do_slice = s->depth <= 8 ? do_hsvhold_slice: do_hsvhold16_slice; |
| 242 | + } |
| 243 | + |
| 244 | + return 0; |
| 245 | +} |
| 246 | + |
| 247 | +static av_cold int query_formats(AVFilterContext *avctx) |
| 248 | +{ |
| 249 | + static const enum AVPixelFormat pixel_fmts[] = { |
| 250 | + AV_PIX_FMT_YUVA420P, |
| 251 | + AV_PIX_FMT_YUVA422P, |
| 252 | + AV_PIX_FMT_YUVA444P, |
| 253 | + AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9, |
| 254 | + AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10, |
| 255 | + AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12, |
| 256 | + AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16, |
| 257 | + AV_PIX_FMT_NONE |
| 258 | + }; |
| 259 | + |
| 260 | + static const enum AVPixelFormat hold_pixel_fmts[] = { |
| 261 | + AV_PIX_FMT_YUV420P, |
| 262 | + AV_PIX_FMT_YUV422P, |
| 263 | + AV_PIX_FMT_YUV444P, |
| 264 | + AV_PIX_FMT_YUVA420P, |
| 265 | + AV_PIX_FMT_YUVA422P, |
| 266 | + AV_PIX_FMT_YUVA444P, |
| 267 | + AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9, |
| 268 | + AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10, |
| 269 | + AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV420P12, |
| 270 | + AV_PIX_FMT_YUV444P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV420P14, |
| 271 | + AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16, |
| 272 | + AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9, |
| 273 | + AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10, |
| 274 | + AV_PIX_FMT_YUVA422P12, AV_PIX_FMT_YUVA444P12, |
| 275 | + AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16, |
| 276 | + AV_PIX_FMT_NONE |
| 277 | + }; |
| 278 | + const enum AVPixelFormat *pix_fmts; |
| 279 | + |
| 280 | + pix_fmts = !strcmp(avctx->filter->name, "hsvhold") ? hold_pixel_fmts : pixel_fmts; |
| 281 | + |
| 282 | + return ff_set_common_formats_from_list(avctx, pix_fmts); |
| 283 | +} |
| 284 | + |
| 285 | +static av_cold int config_input(AVFilterLink *inlink) |
| 286 | +{ |
| 287 | + AVFilterContext *avctx = inlink->dst; |
| 288 | + HSVKeyContext *s = avctx->priv; |
| 289 | + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format); |
| 290 | + |
| 291 | + s->hsub_log2 = desc->log2_chroma_w; |
| 292 | + s->vsub_log2 = desc->log2_chroma_h; |
| 293 | + |
| 294 | + return 0; |
| 295 | +} |
| 296 | + |
| 297 | +static const AVFilterPad hsvkey_inputs[] = { |
| 298 | + { |
| 299 | + .name = "default", |
| 300 | + .type = AVMEDIA_TYPE_VIDEO, |
| 301 | + .flags = AVFILTERPAD_FLAG_NEEDS_WRITABLE, |
| 302 | + .filter_frame = filter_frame, |
| 303 | + .config_props = config_input, |
| 304 | + }, |
| 305 | +}; |
| 306 | + |
| 307 | +static const AVFilterPad hsvkey_outputs[] = { |
| 308 | + { |
| 309 | + .name = "default", |
| 310 | + .type = AVMEDIA_TYPE_VIDEO, |
| 311 | + .config_props = config_output, |
| 312 | + }, |
| 313 | +}; |
| 314 | + |
| 315 | +#define OFFSET(x) offsetof(HSVKeyContext, x) |
| 316 | +#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_RUNTIME_PARAM |
| 317 | + |
| 318 | +static const AVOption hsvkey_options[] = { |
| 319 | + { "hue", "set the hue value", OFFSET(hue_opt), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -360, 360, FLAGS }, |
| 320 | + { "sat", "set the saturation value", OFFSET(sat), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -1, 1, FLAGS }, |
| 321 | + { "val", "set the value value", OFFSET(val), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -1, 1, FLAGS }, |
| 322 | + { "similarity", "set the hsvkey similarity value", OFFSET(similarity), AV_OPT_TYPE_FLOAT, { .dbl = 0.01}, 0.00001, 1.0, FLAGS }, |
| 323 | + { "blend", "set the hsvkey blend value", OFFSET(blend), AV_OPT_TYPE_FLOAT, { .dbl = 0.0 }, 0.0, 1.0, FLAGS }, |
| 324 | + { NULL } |
| 325 | +}; |
| 326 | + |
| 327 | +AVFILTER_DEFINE_CLASS(hsvkey); |
| 328 | + |
| 329 | +const AVFilter ff_vf_hsvkey = { |
| 330 | + .name = "hsvkey", |
| 331 | + .description = NULL_IF_CONFIG_SMALL("Turns a certain HSV range into transparency. Operates on YUV colors."), |
| 332 | + .priv_size = sizeof(HSVKeyContext), |
| 333 | + .priv_class = &hsvkey_class, |
| 334 | + .query_formats = query_formats, |
| 335 | + FILTER_INPUTS(hsvkey_inputs), |
| 336 | + FILTER_OUTPUTS(hsvkey_outputs), |
| 337 | + .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS, |
| 338 | + .process_command = ff_filter_process_command, |
| 339 | +}; |
| 340 | + |
| 341 | +static const AVOption hsvhold_options[] = { |
| 342 | + { "hue", "set the hue value", OFFSET(hue_opt), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -360, 360, FLAGS }, |
| 343 | + { "sat", "set the saturation value", OFFSET(sat), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -1, 1, FLAGS }, |
| 344 | + { "val", "set the value value", OFFSET(val), AV_OPT_TYPE_FLOAT, { .dbl = 0 }, -1, 1, FLAGS }, |
| 345 | + { "similarity", "set the hsvhold similarity value", OFFSET(similarity), AV_OPT_TYPE_FLOAT, { .dbl = 0.01 }, 0.00001, 1.0, FLAGS }, |
| 346 | + { "blend", "set the hsvhold blend value", OFFSET(blend), AV_OPT_TYPE_FLOAT, { .dbl = 0.0 }, 0.0, 1.0, FLAGS }, |
| 347 | + { NULL } |
| 348 | +}; |
| 349 | + |
| 350 | +static const AVFilterPad hsvhold_inputs[] = { |
| 351 | + { |
| 352 | + .name = "default", |
| 353 | + .type = AVMEDIA_TYPE_VIDEO, |
| 354 | + .flags = AVFILTERPAD_FLAG_NEEDS_WRITABLE, |
| 355 | + .filter_frame = filter_frame, |
| 356 | + .config_props = config_input, |
| 357 | + }, |
| 358 | +}; |
| 359 | + |
| 360 | +static const AVFilterPad hsvhold_outputs[] = { |
| 361 | + { |
| 362 | + .name = "default", |
| 363 | + .type = AVMEDIA_TYPE_VIDEO, |
| 364 | + .config_props = config_output, |
| 365 | + }, |
| 366 | +}; |
| 367 | + |
| 368 | +AVFILTER_DEFINE_CLASS(hsvhold); |
| 369 | + |
| 370 | +const AVFilter ff_vf_hsvhold = { |
| 371 | + .name = "hsvhold", |
| 372 | + .description = NULL_IF_CONFIG_SMALL("Turns a certain HSV range into gray."), |
| 373 | + .priv_size = sizeof(HSVKeyContext), |
| 374 | + .priv_class = &hsvhold_class, |
| 375 | + .query_formats = query_formats, |
| 376 | + FILTER_INPUTS(hsvhold_inputs), |
| 377 | + FILTER_OUTPUTS(hsvhold_outputs), |
| 378 | + .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS, |
| 379 | + .process_command = ff_filter_process_command, |
| 380 | +}; |