FFmpeg coverage


Directory: ../../../ffmpeg/
File: src/libavfilter/vf_dnn_detect.c
Date: 2024-04-19 07:31:02
Exec Total Coverage
Lines: 0 478 0.0%
Functions: 0 21 0.0%
Branches: 0 241 0.0%

Line Branch Exec Source
1 /*
2 * This file is part of FFmpeg.
3 *
4 * FFmpeg is free software; you can redistribute it and/or
5 * modify it under the terms of the GNU Lesser General Public
6 * License as published by the Free Software Foundation; either
7 * version 2.1 of the License, or (at your option) any later version.
8 *
9 * FFmpeg is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
12 * Lesser General Public License for more details.
13 *
14 * You should have received a copy of the GNU Lesser General Public
15 * License along with FFmpeg; if not, write to the Free Software
16 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17 */
18
19 /**
20 * @file
21 * implementing an object detecting filter using deep learning networks.
22 */
23
24 #include "libavutil/file_open.h"
25 #include "libavutil/mem.h"
26 #include "libavutil/opt.h"
27 #include "filters.h"
28 #include "dnn_filter_common.h"
29 #include "internal.h"
30 #include "video.h"
31 #include "libavutil/time.h"
32 #include "libavutil/avstring.h"
33 #include "libavutil/detection_bbox.h"
34 #include "libavutil/fifo.h"
35
36 typedef enum {
37 DDMT_SSD,
38 DDMT_YOLOV1V2,
39 DDMT_YOLOV3,
40 DDMT_YOLOV4
41 } DNNDetectionModelType;
42
43 typedef struct DnnDetectContext {
44 const AVClass *class;
45 DnnContext dnnctx;
46 float confidence;
47 char *labels_filename;
48 char **labels;
49 int label_count;
50 DNNDetectionModelType model_type;
51 int cell_w;
52 int cell_h;
53 int nb_classes;
54 AVFifo *bboxes_fifo;
55 int scale_width;
56 int scale_height;
57 char *anchors_str;
58 float *anchors;
59 int nb_anchor;
60 } DnnDetectContext;
61
62 #define OFFSET(x) offsetof(DnnDetectContext, dnnctx.x)
63 #define OFFSET2(x) offsetof(DnnDetectContext, x)
64 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM
65 static const AVOption dnn_detect_options[] = {
66 { "dnn_backend", "DNN backend", OFFSET(backend_type), AV_OPT_TYPE_INT, { .i64 = DNN_OV }, INT_MIN, INT_MAX, FLAGS, .unit = "backend" },
67 #if (CONFIG_LIBTENSORFLOW == 1)
68 { "tensorflow", "tensorflow backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = DNN_TF }, 0, 0, FLAGS, .unit = "backend" },
69 #endif
70 #if (CONFIG_LIBOPENVINO == 1)
71 { "openvino", "openvino backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = DNN_OV }, 0, 0, FLAGS, .unit = "backend" },
72 #endif
73 DNN_COMMON_OPTIONS
74 { "confidence", "threshold of confidence", OFFSET2(confidence), AV_OPT_TYPE_FLOAT, { .dbl = 0.5 }, 0, 1, FLAGS},
75 { "labels", "path to labels file", OFFSET2(labels_filename), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS },
76 { "model_type", "DNN detection model type", OFFSET2(model_type), AV_OPT_TYPE_INT, { .i64 = DDMT_SSD }, INT_MIN, INT_MAX, FLAGS, .unit = "model_type" },
77 { "ssd", "output shape [1, 1, N, 7]", 0, AV_OPT_TYPE_CONST, { .i64 = DDMT_SSD }, 0, 0, FLAGS, .unit = "model_type" },
78 { "yolo", "output shape [1, N*Cx*Cy*DetectionBox]", 0, AV_OPT_TYPE_CONST, { .i64 = DDMT_YOLOV1V2 }, 0, 0, FLAGS, .unit = "model_type" },
79 { "yolov3", "outputs shape [1, N*D, Cx, Cy]", 0, AV_OPT_TYPE_CONST, { .i64 = DDMT_YOLOV3 }, 0, 0, FLAGS, .unit = "model_type" },
80 { "yolov4", "outputs shape [1, N*D, Cx, Cy]", 0, AV_OPT_TYPE_CONST, { .i64 = DDMT_YOLOV4 }, 0, 0, FLAGS, .unit = "model_type" },
81 { "cell_w", "cell width", OFFSET2(cell_w), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INTMAX_MAX, FLAGS },
82 { "cell_h", "cell height", OFFSET2(cell_h), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INTMAX_MAX, FLAGS },
83 { "nb_classes", "The number of class", OFFSET2(nb_classes), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INTMAX_MAX, FLAGS },
84 { "anchors", "anchors, splited by '&'", OFFSET2(anchors_str), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS },
85 { NULL }
86 };
87
88 AVFILTER_DEFINE_CLASS(dnn_detect);
89
90 static inline float sigmoid(float x) {
91 return 1.f / (1.f + exp(-x));
92 }
93
94 static inline float linear(float x) {
95 return x;
96 }
97
98 static int dnn_detect_get_label_id(int nb_classes, int cell_size, float *label_data)
99 {
100 float max_prob = 0;
101 int label_id = 0;
102 for (int i = 0; i < nb_classes; i++) {
103 if (label_data[i * cell_size] > max_prob) {
104 max_prob = label_data[i * cell_size];
105 label_id = i;
106 }
107 }
108 return label_id;
109 }
110
111 static int dnn_detect_parse_anchors(char *anchors_str, float **anchors)
112 {
113 char *saveptr = NULL, *token;
114 float *anchors_buf;
115 int nb_anchor = 0, i = 0;
116 while(anchors_str[i] != '\0') {
117 if(anchors_str[i] == '&')
118 nb_anchor++;
119 i++;
120 }
121 nb_anchor++;
122 anchors_buf = av_mallocz(nb_anchor * sizeof(**anchors));
123 if (!anchors_buf) {
124 return 0;
125 }
126 for (int i = 0; i < nb_anchor; i++) {
127 token = av_strtok(anchors_str, "&", &saveptr);
128 if (!token) {
129 av_freep(&anchors_buf);
130 return 0;
131 }
132 anchors_buf[i] = strtof(token, NULL);
133 anchors_str = NULL;
134 }
135 *anchors = anchors_buf;
136 return nb_anchor;
137 }
138
139 /* Calculate Intersection Over Union */
140 static float dnn_detect_IOU(AVDetectionBBox *bbox1, AVDetectionBBox *bbox2)
141 {
142 float overlapping_width = FFMIN(bbox1->x + bbox1->w, bbox2->x + bbox2->w) - FFMAX(bbox1->x, bbox2->x);
143 float overlapping_height = FFMIN(bbox1->y + bbox1->h, bbox2->y + bbox2->h) - FFMAX(bbox1->y, bbox2->y);
144 float intersection_area =
145 (overlapping_width < 0 || overlapping_height < 0) ? 0 : overlapping_height * overlapping_width;
146 float union_area = bbox1->w * bbox1->h + bbox2->w * bbox2->h - intersection_area;
147 return intersection_area / union_area;
148 }
149
150 static int dnn_detect_parse_yolo_output(AVFrame *frame, DNNData *output, int output_index,
151 AVFilterContext *filter_ctx)
152 {
153 DnnDetectContext *ctx = filter_ctx->priv;
154 float conf_threshold = ctx->confidence;
155 int detection_boxes, box_size;
156 int cell_w = 0, cell_h = 0, scale_w = 0, scale_h = 0;
157 int nb_classes = ctx->nb_classes;
158 float *output_data = output[output_index].data;
159 float *anchors = ctx->anchors;
160 AVDetectionBBox *bbox;
161 float (*post_process_raw_data)(float x) = linear;
162 int is_NHWC = 0;
163
164 if (ctx->model_type == DDMT_YOLOV1V2) {
165 cell_w = ctx->cell_w;
166 cell_h = ctx->cell_h;
167 scale_w = cell_w;
168 scale_h = cell_h;
169 } else {
170 if (output[output_index].dims[2] != output[output_index].dims[3] &&
171 output[output_index].dims[2] == output[output_index].dims[1]) {
172 is_NHWC = 1;
173 cell_w = output[output_index].dims[2];
174 cell_h = output[output_index].dims[1];
175 } else {
176 cell_w = output[output_index].dims[3];
177 cell_h = output[output_index].dims[2];
178 }
179 scale_w = ctx->scale_width;
180 scale_h = ctx->scale_height;
181 }
182 box_size = nb_classes + 5;
183
184 switch (ctx->model_type) {
185 case DDMT_YOLOV1V2:
186 case DDMT_YOLOV3:
187 post_process_raw_data = linear;
188 break;
189 case DDMT_YOLOV4:
190 post_process_raw_data = sigmoid;
191 break;
192 }
193
194 if (!cell_h || !cell_w) {
195 av_log(filter_ctx, AV_LOG_ERROR, "cell_w and cell_h are detected\n");
196 return AVERROR(EINVAL);
197 }
198
199 if (!nb_classes) {
200 av_log(filter_ctx, AV_LOG_ERROR, "nb_classes is not set\n");
201 return AVERROR(EINVAL);
202 }
203
204 if (!anchors) {
205 av_log(filter_ctx, AV_LOG_ERROR, "anchors is not set\n");
206 return AVERROR(EINVAL);
207 }
208
209 if (output[output_index].dims[1] * output[output_index].dims[2] *
210 output[output_index].dims[3] % (box_size * cell_w * cell_h)) {
211 av_log(filter_ctx, AV_LOG_ERROR, "wrong cell_w, cell_h or nb_classes\n");
212 return AVERROR(EINVAL);
213 }
214 detection_boxes = output[output_index].dims[1] *
215 output[output_index].dims[2] *
216 output[output_index].dims[3] / box_size / cell_w / cell_h;
217
218 anchors = anchors + (detection_boxes * output_index * 2);
219 /**
220 * find all candidate bbox
221 * yolo output can be reshaped to [B, N*D, Cx, Cy]
222 * Detection box 'D' has format [`x`, `y`, `h`, `w`, `box_score`, `class_no_1`, ...,]
223 **/
224 for (int box_id = 0; box_id < detection_boxes; box_id++) {
225 for (int cx = 0; cx < cell_w; cx++)
226 for (int cy = 0; cy < cell_h; cy++) {
227 float x, y, w, h, conf;
228 float *detection_boxes_data;
229 int label_id;
230
231 if (is_NHWC) {
232 detection_boxes_data = output_data +
233 ((cy * cell_w + cx) * detection_boxes + box_id) * box_size;
234 conf = post_process_raw_data(detection_boxes_data[4]);
235 } else {
236 detection_boxes_data = output_data + box_id * box_size * cell_w * cell_h;
237 conf = post_process_raw_data(
238 detection_boxes_data[cy * cell_w + cx + 4 * cell_w * cell_h]);
239 }
240
241 if (is_NHWC) {
242 x = post_process_raw_data(detection_boxes_data[0]);
243 y = post_process_raw_data(detection_boxes_data[1]);
244 w = detection_boxes_data[2];
245 h = detection_boxes_data[3];
246 label_id = dnn_detect_get_label_id(ctx->nb_classes, 1, detection_boxes_data + 5);
247 conf = conf * post_process_raw_data(detection_boxes_data[label_id + 5]);
248 } else {
249 x = post_process_raw_data(detection_boxes_data[cy * cell_w + cx]);
250 y = post_process_raw_data(detection_boxes_data[cy * cell_w + cx + cell_w * cell_h]);
251 w = detection_boxes_data[cy * cell_w + cx + 2 * cell_w * cell_h];
252 h = detection_boxes_data[cy * cell_w + cx + 3 * cell_w * cell_h];
253 label_id = dnn_detect_get_label_id(ctx->nb_classes, cell_w * cell_h,
254 detection_boxes_data + cy * cell_w + cx + 5 * cell_w * cell_h);
255 conf = conf * post_process_raw_data(
256 detection_boxes_data[cy * cell_w + cx + (label_id + 5) * cell_w * cell_h]);
257 }
258 if (conf < conf_threshold) {
259 continue;
260 }
261
262 bbox = av_mallocz(sizeof(*bbox));
263 if (!bbox)
264 return AVERROR(ENOMEM);
265
266 bbox->w = exp(w) * anchors[box_id * 2] * frame->width / scale_w;
267 bbox->h = exp(h) * anchors[box_id * 2 + 1] * frame->height / scale_h;
268 bbox->x = (cx + x) / cell_w * frame->width - bbox->w / 2;
269 bbox->y = (cy + y) / cell_h * frame->height - bbox->h / 2;
270 bbox->detect_confidence = av_make_q((int)(conf * 10000), 10000);
271 if (ctx->labels && label_id < ctx->label_count) {
272 av_strlcpy(bbox->detect_label, ctx->labels[label_id], sizeof(bbox->detect_label));
273 } else {
274 snprintf(bbox->detect_label, sizeof(bbox->detect_label), "%d", label_id);
275 }
276
277 if (av_fifo_write(ctx->bboxes_fifo, &bbox, 1) < 0) {
278 av_freep(&bbox);
279 return AVERROR(ENOMEM);
280 }
281 bbox = NULL;
282 }
283 }
284 return 0;
285 }
286
287 static int dnn_detect_fill_side_data(AVFrame *frame, AVFilterContext *filter_ctx)
288 {
289 DnnDetectContext *ctx = filter_ctx->priv;
290 float conf_threshold = ctx->confidence;
291 AVDetectionBBox *bbox;
292 int nb_bboxes = 0;
293 AVDetectionBBoxHeader *header;
294 if (av_fifo_can_read(ctx->bboxes_fifo) == 0) {
295 av_log(filter_ctx, AV_LOG_VERBOSE, "nothing detected in this frame.\n");
296 return 0;
297 }
298
299 /* remove overlap bboxes */
300 for (int i = 0; i < av_fifo_can_read(ctx->bboxes_fifo); i++){
301 av_fifo_peek(ctx->bboxes_fifo, &bbox, 1, i);
302 for (int j = 0; j < av_fifo_can_read(ctx->bboxes_fifo); j++) {
303 AVDetectionBBox *overlap_bbox;
304 av_fifo_peek(ctx->bboxes_fifo, &overlap_bbox, 1, j);
305 if (!strcmp(bbox->detect_label, overlap_bbox->detect_label) &&
306 av_cmp_q(bbox->detect_confidence, overlap_bbox->detect_confidence) < 0 &&
307 dnn_detect_IOU(bbox, overlap_bbox) >= conf_threshold) {
308 bbox->classify_count = -1; // bad result
309 nb_bboxes++;
310 break;
311 }
312 }
313 }
314 nb_bboxes = av_fifo_can_read(ctx->bboxes_fifo) - nb_bboxes;
315 header = av_detection_bbox_create_side_data(frame, nb_bboxes);
316 if (!header) {
317 av_log(filter_ctx, AV_LOG_ERROR, "failed to create side data with %d bounding boxes\n", nb_bboxes);
318 return -1;
319 }
320 av_strlcpy(header->source, ctx->dnnctx.model_filename, sizeof(header->source));
321
322 while(av_fifo_can_read(ctx->bboxes_fifo)) {
323 AVDetectionBBox *candidate_bbox;
324 av_fifo_read(ctx->bboxes_fifo, &candidate_bbox, 1);
325
326 if (nb_bboxes > 0 && candidate_bbox->classify_count != -1) {
327 bbox = av_get_detection_bbox(header, header->nb_bboxes - nb_bboxes);
328 memcpy(bbox, candidate_bbox, sizeof(*bbox));
329 nb_bboxes--;
330 }
331 av_freep(&candidate_bbox);
332 }
333 return 0;
334 }
335
336 static int dnn_detect_post_proc_yolo(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx)
337 {
338 int ret = 0;
339 ret = dnn_detect_parse_yolo_output(frame, output, 0, filter_ctx);
340 if (ret < 0)
341 return ret;
342 ret = dnn_detect_fill_side_data(frame, filter_ctx);
343 if (ret < 0)
344 return ret;
345 return 0;
346 }
347
348 static int dnn_detect_post_proc_yolov3(AVFrame *frame, DNNData *output,
349 AVFilterContext *filter_ctx, int nb_outputs)
350 {
351 int ret = 0;
352 for (int i = 0; i < nb_outputs; i++) {
353 ret = dnn_detect_parse_yolo_output(frame, output, i, filter_ctx);
354 if (ret < 0)
355 return ret;
356 }
357 ret = dnn_detect_fill_side_data(frame, filter_ctx);
358 if (ret < 0)
359 return ret;
360 return 0;
361 }
362
363 static int dnn_detect_post_proc_ssd(AVFrame *frame, DNNData *output, int nb_outputs,
364 AVFilterContext *filter_ctx)
365 {
366 DnnDetectContext *ctx = filter_ctx->priv;
367 float conf_threshold = ctx->confidence;
368 int proposal_count = 0;
369 int detect_size = 0;
370 float *detections = NULL, *labels = NULL;
371 int nb_bboxes = 0;
372 AVDetectionBBoxHeader *header;
373 AVDetectionBBox *bbox;
374 int scale_w = ctx->scale_width;
375 int scale_h = ctx->scale_height;
376
377 if (nb_outputs == 1 && output->dims[3] == 7) {
378 proposal_count = output->dims[2];
379 detect_size = output->dims[3];
380 detections = output->data;
381 } else if (nb_outputs == 2 && output[0].dims[3] == 5) {
382 proposal_count = output[0].dims[2];
383 detect_size = output[0].dims[3];
384 detections = output[0].data;
385 labels = output[1].data;
386 } else if (nb_outputs == 2 && output[1].dims[3] == 5) {
387 proposal_count = output[1].dims[2];
388 detect_size = output[1].dims[3];
389 detections = output[1].data;
390 labels = output[0].data;
391 } else {
392 av_log(filter_ctx, AV_LOG_ERROR, "Model output shape doesn't match ssd requirement.\n");
393 return AVERROR(EINVAL);
394 }
395
396 if (proposal_count == 0)
397 return 0;
398
399 for (int i = 0; i < proposal_count; ++i) {
400 float conf;
401 if (nb_outputs == 1)
402 conf = detections[i * detect_size + 2];
403 else
404 conf = detections[i * detect_size + 4];
405 if (conf < conf_threshold) {
406 continue;
407 }
408 nb_bboxes++;
409 }
410
411 if (nb_bboxes == 0) {
412 av_log(filter_ctx, AV_LOG_VERBOSE, "nothing detected in this frame.\n");
413 return 0;
414 }
415
416 header = av_detection_bbox_create_side_data(frame, nb_bboxes);
417 if (!header) {
418 av_log(filter_ctx, AV_LOG_ERROR, "failed to create side data with %d bounding boxes\n", nb_bboxes);
419 return -1;
420 }
421
422 av_strlcpy(header->source, ctx->dnnctx.model_filename, sizeof(header->source));
423
424 for (int i = 0; i < proposal_count; ++i) {
425 int av_unused image_id = (int)detections[i * detect_size + 0];
426 int label_id;
427 float conf, x0, y0, x1, y1;
428
429 if (nb_outputs == 1) {
430 label_id = (int)detections[i * detect_size + 1];
431 conf = detections[i * detect_size + 2];
432 x0 = detections[i * detect_size + 3];
433 y0 = detections[i * detect_size + 4];
434 x1 = detections[i * detect_size + 5];
435 y1 = detections[i * detect_size + 6];
436 } else {
437 label_id = (int)labels[i];
438 x0 = detections[i * detect_size] / scale_w;
439 y0 = detections[i * detect_size + 1] / scale_h;
440 x1 = detections[i * detect_size + 2] / scale_w;
441 y1 = detections[i * detect_size + 3] / scale_h;
442 conf = detections[i * detect_size + 4];
443 }
444
445 if (conf < conf_threshold) {
446 continue;
447 }
448
449 bbox = av_get_detection_bbox(header, header->nb_bboxes - nb_bboxes);
450 bbox->x = (int)(x0 * frame->width);
451 bbox->w = (int)(x1 * frame->width) - bbox->x;
452 bbox->y = (int)(y0 * frame->height);
453 bbox->h = (int)(y1 * frame->height) - bbox->y;
454
455 bbox->detect_confidence = av_make_q((int)(conf * 10000), 10000);
456 bbox->classify_count = 0;
457
458 if (ctx->labels && label_id < ctx->label_count) {
459 av_strlcpy(bbox->detect_label, ctx->labels[label_id], sizeof(bbox->detect_label));
460 } else {
461 snprintf(bbox->detect_label, sizeof(bbox->detect_label), "%d", label_id);
462 }
463
464 nb_bboxes--;
465 if (nb_bboxes == 0) {
466 break;
467 }
468 }
469 return 0;
470 }
471
472 static int dnn_detect_post_proc_ov(AVFrame *frame, DNNData *output, int nb_outputs,
473 AVFilterContext *filter_ctx)
474 {
475 AVFrameSideData *sd;
476 DnnDetectContext *ctx = filter_ctx->priv;
477 int ret = 0;
478
479 sd = av_frame_get_side_data(frame, AV_FRAME_DATA_DETECTION_BBOXES);
480 if (sd) {
481 av_log(filter_ctx, AV_LOG_ERROR, "already have bounding boxes in side data.\n");
482 return -1;
483 }
484
485 switch (ctx->model_type) {
486 case DDMT_SSD:
487 ret = dnn_detect_post_proc_ssd(frame, output, nb_outputs, filter_ctx);
488 if (ret < 0)
489 return ret;
490 break;
491 case DDMT_YOLOV1V2:
492 ret = dnn_detect_post_proc_yolo(frame, output, filter_ctx);
493 if (ret < 0)
494 return ret;
495 break;
496 case DDMT_YOLOV3:
497 case DDMT_YOLOV4:
498 ret = dnn_detect_post_proc_yolov3(frame, output, filter_ctx, nb_outputs);
499 if (ret < 0)
500 return ret;
501 break;
502 }
503 return 0;
504 }
505
506 static int dnn_detect_post_proc_tf(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx)
507 {
508 DnnDetectContext *ctx = filter_ctx->priv;
509 int proposal_count;
510 float conf_threshold = ctx->confidence;
511 float *conf, *position, *label_id, x0, y0, x1, y1;
512 int nb_bboxes = 0;
513 AVFrameSideData *sd;
514 AVDetectionBBox *bbox;
515 AVDetectionBBoxHeader *header;
516
517 proposal_count = *(float *)(output[0].data);
518 conf = output[1].data;
519 position = output[3].data;
520 label_id = output[2].data;
521
522 sd = av_frame_get_side_data(frame, AV_FRAME_DATA_DETECTION_BBOXES);
523 if (sd) {
524 av_log(filter_ctx, AV_LOG_ERROR, "already have dnn bounding boxes in side data.\n");
525 return -1;
526 }
527
528 for (int i = 0; i < proposal_count; ++i) {
529 if (conf[i] < conf_threshold)
530 continue;
531 nb_bboxes++;
532 }
533
534 if (nb_bboxes == 0) {
535 av_log(filter_ctx, AV_LOG_VERBOSE, "nothing detected in this frame.\n");
536 return 0;
537 }
538
539 header = av_detection_bbox_create_side_data(frame, nb_bboxes);
540 if (!header) {
541 av_log(filter_ctx, AV_LOG_ERROR, "failed to create side data with %d bounding boxes\n", nb_bboxes);
542 return -1;
543 }
544
545 av_strlcpy(header->source, ctx->dnnctx.model_filename, sizeof(header->source));
546
547 for (int i = 0; i < proposal_count; ++i) {
548 y0 = position[i * 4];
549 x0 = position[i * 4 + 1];
550 y1 = position[i * 4 + 2];
551 x1 = position[i * 4 + 3];
552
553 bbox = av_get_detection_bbox(header, i);
554
555 if (conf[i] < conf_threshold) {
556 continue;
557 }
558
559 bbox->x = (int)(x0 * frame->width);
560 bbox->w = (int)(x1 * frame->width) - bbox->x;
561 bbox->y = (int)(y0 * frame->height);
562 bbox->h = (int)(y1 * frame->height) - bbox->y;
563
564 bbox->detect_confidence = av_make_q((int)(conf[i] * 10000), 10000);
565 bbox->classify_count = 0;
566
567 if (ctx->labels && label_id[i] < ctx->label_count) {
568 av_strlcpy(bbox->detect_label, ctx->labels[(int)label_id[i]], sizeof(bbox->detect_label));
569 } else {
570 snprintf(bbox->detect_label, sizeof(bbox->detect_label), "%d", (int)label_id[i]);
571 }
572
573 nb_bboxes--;
574 if (nb_bboxes == 0) {
575 break;
576 }
577 }
578 return 0;
579 }
580
581 static int dnn_detect_post_proc(AVFrame *frame, DNNData *output, uint32_t nb, AVFilterContext *filter_ctx)
582 {
583 DnnDetectContext *ctx = filter_ctx->priv;
584 DnnContext *dnn_ctx = &ctx->dnnctx;
585 switch (dnn_ctx->backend_type) {
586 case DNN_OV:
587 return dnn_detect_post_proc_ov(frame, output, nb, filter_ctx);
588 case DNN_TF:
589 return dnn_detect_post_proc_tf(frame, output, filter_ctx);
590 default:
591 avpriv_report_missing_feature(filter_ctx, "Current dnn backend does not support detect filter\n");
592 return AVERROR(EINVAL);
593 }
594 }
595
596 static void free_detect_labels(DnnDetectContext *ctx)
597 {
598 for (int i = 0; i < ctx->label_count; i++) {
599 av_freep(&ctx->labels[i]);
600 }
601 ctx->label_count = 0;
602 av_freep(&ctx->labels);
603 }
604
605 static int read_detect_label_file(AVFilterContext *context)
606 {
607 int line_len;
608 FILE *file;
609 DnnDetectContext *ctx = context->priv;
610
611 file = avpriv_fopen_utf8(ctx->labels_filename, "r");
612 if (!file){
613 av_log(context, AV_LOG_ERROR, "failed to open file %s\n", ctx->labels_filename);
614 return AVERROR(EINVAL);
615 }
616
617 while (!feof(file)) {
618 char *label;
619 char buf[256];
620 if (!fgets(buf, 256, file)) {
621 break;
622 }
623
624 line_len = strlen(buf);
625 while (line_len) {
626 int i = line_len - 1;
627 if (buf[i] == '\n' || buf[i] == '\r' || buf[i] == ' ') {
628 buf[i] = '\0';
629 line_len--;
630 } else {
631 break;
632 }
633 }
634
635 if (line_len == 0) // empty line
636 continue;
637
638 if (line_len >= AV_DETECTION_BBOX_LABEL_NAME_MAX_SIZE) {
639 av_log(context, AV_LOG_ERROR, "label %s too long\n", buf);
640 fclose(file);
641 return AVERROR(EINVAL);
642 }
643
644 label = av_strdup(buf);
645 if (!label) {
646 av_log(context, AV_LOG_ERROR, "failed to allocate memory for label %s\n", buf);
647 fclose(file);
648 return AVERROR(ENOMEM);
649 }
650
651 if (av_dynarray_add_nofree(&ctx->labels, &ctx->label_count, label) < 0) {
652 av_log(context, AV_LOG_ERROR, "failed to do av_dynarray_add\n");
653 fclose(file);
654 av_freep(&label);
655 return AVERROR(ENOMEM);
656 }
657 }
658
659 fclose(file);
660 return 0;
661 }
662
663 static int check_output_nb(DnnDetectContext *ctx, DNNBackendType backend_type, int output_nb)
664 {
665 switch(backend_type) {
666 case DNN_TF:
667 if (output_nb != 4) {
668 av_log(ctx, AV_LOG_ERROR, "Only support tensorflow detect model with 4 outputs, \
669 but get %d instead\n", output_nb);
670 return AVERROR(EINVAL);
671 }
672 return 0;
673 case DNN_OV:
674 return 0;
675 default:
676 avpriv_report_missing_feature(ctx, "Dnn detect filter does not support current backend\n");
677 return AVERROR(EINVAL);
678 }
679 return 0;
680 }
681
682 static av_cold int dnn_detect_init(AVFilterContext *context)
683 {
684 DnnDetectContext *ctx = context->priv;
685 DnnContext *dnn_ctx = &ctx->dnnctx;
686 int ret;
687
688 ret = ff_dnn_init(&ctx->dnnctx, DFT_ANALYTICS_DETECT, context);
689 if (ret < 0)
690 return ret;
691 ret = check_output_nb(ctx, dnn_ctx->backend_type, dnn_ctx->nb_outputs);
692 if (ret < 0)
693 return ret;
694 ctx->bboxes_fifo = av_fifo_alloc2(1, sizeof(AVDetectionBBox *), AV_FIFO_FLAG_AUTO_GROW);
695 if (!ctx->bboxes_fifo)
696 return AVERROR(ENOMEM);
697 ff_dnn_set_detect_post_proc(&ctx->dnnctx, dnn_detect_post_proc);
698
699 if (ctx->labels_filename) {
700 return read_detect_label_file(context);
701 }
702 if (ctx->anchors_str) {
703 ret = dnn_detect_parse_anchors(ctx->anchors_str, &ctx->anchors);
704 if (!ctx->anchors) {
705 av_log(context, AV_LOG_ERROR, "failed to parse anchors_str\n");
706 return AVERROR(EINVAL);
707 }
708 ctx->nb_anchor = ret;
709 }
710 return 0;
711 }
712
713 static const enum AVPixelFormat pix_fmts[] = {
714 AV_PIX_FMT_RGB24, AV_PIX_FMT_BGR24,
715 AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAYF32,
716 AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P,
717 AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P,
718 AV_PIX_FMT_NV12,
719 AV_PIX_FMT_NONE
720 };
721
722 static int dnn_detect_flush_frame(AVFilterLink *outlink, int64_t pts, int64_t *out_pts)
723 {
724 DnnDetectContext *ctx = outlink->src->priv;
725 int ret;
726 DNNAsyncStatusType async_state;
727
728 ret = ff_dnn_flush(&ctx->dnnctx);
729 if (ret != 0) {
730 return -1;
731 }
732
733 do {
734 AVFrame *in_frame = NULL;
735 AVFrame *out_frame = NULL;
736 async_state = ff_dnn_get_result(&ctx->dnnctx, &in_frame, &out_frame);
737 if (async_state == DAST_SUCCESS) {
738 ret = ff_filter_frame(outlink, in_frame);
739 if (ret < 0)
740 return ret;
741 if (out_pts)
742 *out_pts = in_frame->pts + pts;
743 }
744 av_usleep(5000);
745 } while (async_state >= DAST_NOT_READY);
746
747 return 0;
748 }
749
750 static int dnn_detect_activate(AVFilterContext *filter_ctx)
751 {
752 AVFilterLink *inlink = filter_ctx->inputs[0];
753 AVFilterLink *outlink = filter_ctx->outputs[0];
754 DnnDetectContext *ctx = filter_ctx->priv;
755 AVFrame *in = NULL;
756 int64_t pts;
757 int ret, status;
758 int got_frame = 0;
759 int async_state;
760
761 FF_FILTER_FORWARD_STATUS_BACK(outlink, inlink);
762
763 do {
764 // drain all input frames
765 ret = ff_inlink_consume_frame(inlink, &in);
766 if (ret < 0)
767 return ret;
768 if (ret > 0) {
769 if (ff_dnn_execute_model(&ctx->dnnctx, in, NULL) != 0) {
770 return AVERROR(EIO);
771 }
772 }
773 } while (ret > 0);
774
775 // drain all processed frames
776 do {
777 AVFrame *in_frame = NULL;
778 AVFrame *out_frame = NULL;
779 async_state = ff_dnn_get_result(&ctx->dnnctx, &in_frame, &out_frame);
780 if (async_state == DAST_SUCCESS) {
781 ret = ff_filter_frame(outlink, in_frame);
782 if (ret < 0)
783 return ret;
784 got_frame = 1;
785 }
786 } while (async_state == DAST_SUCCESS);
787
788 // if frame got, schedule to next filter
789 if (got_frame)
790 return 0;
791
792 if (ff_inlink_acknowledge_status(inlink, &status, &pts)) {
793 if (status == AVERROR_EOF) {
794 int64_t out_pts = pts;
795 ret = dnn_detect_flush_frame(outlink, pts, &out_pts);
796 ff_outlink_set_status(outlink, status, out_pts);
797 return ret;
798 }
799 }
800
801 FF_FILTER_FORWARD_WANTED(outlink, inlink);
802
803 return 0;
804 }
805
806 static av_cold void dnn_detect_uninit(AVFilterContext *context)
807 {
808 DnnDetectContext *ctx = context->priv;
809 AVDetectionBBox *bbox;
810 ff_dnn_uninit(&ctx->dnnctx);
811 while(av_fifo_can_read(ctx->bboxes_fifo)) {
812 av_fifo_read(ctx->bboxes_fifo, &bbox, 1);
813 av_freep(&bbox);
814 }
815 av_fifo_freep2(&ctx->bboxes_fifo);
816 av_freep(&ctx->anchors);
817 free_detect_labels(ctx);
818 }
819
820 static int config_input(AVFilterLink *inlink)
821 {
822 AVFilterContext *context = inlink->dst;
823 DnnDetectContext *ctx = context->priv;
824 DNNData model_input;
825 int ret, width_idx, height_idx;
826
827 ret = ff_dnn_get_input(&ctx->dnnctx, &model_input);
828 if (ret != 0) {
829 av_log(ctx, AV_LOG_ERROR, "could not get input from the model\n");
830 return ret;
831 }
832 width_idx = dnn_get_width_idx_by_layout(model_input.layout);
833 height_idx = dnn_get_height_idx_by_layout(model_input.layout);
834 ctx->scale_width = model_input.dims[width_idx] == -1 ? inlink->w :
835 model_input.dims[width_idx];
836 ctx->scale_height = model_input.dims[height_idx] == -1 ? inlink->h :
837 model_input.dims[height_idx];
838
839 return 0;
840 }
841
842 static const AVFilterPad dnn_detect_inputs[] = {
843 {
844 .name = "default",
845 .type = AVMEDIA_TYPE_VIDEO,
846 .config_props = config_input,
847 },
848 };
849
850 const AVFilter ff_vf_dnn_detect = {
851 .name = "dnn_detect",
852 .description = NULL_IF_CONFIG_SMALL("Apply DNN detect filter to the input."),
853 .priv_size = sizeof(DnnDetectContext),
854 .init = dnn_detect_init,
855 .uninit = dnn_detect_uninit,
856 FILTER_INPUTS(dnn_detect_inputs),
857 FILTER_OUTPUTS(ff_video_default_filterpad),
858 FILTER_PIXFMTS_ARRAY(pix_fmts),
859 .priv_class = &dnn_detect_class,
860 .activate = dnn_detect_activate,
861 };
862