There are many kinds of detection DNN model and they have different preprocess and postprocess methods. To support more models, "model_type" option is added to help to choose preprocess and postprocess function.
Unified · Split
| @@ -31,6 +31,10 @@ · 30 unmodified lines | |||
| 31 | #include "libavutil/avstring.h" | 31 | #include "libavutil/avstring.h" |
| 32 | #include "libavutil/detection_bbox.h" | 32 | #include "libavutil/detection_bbox.h" |
| 33 | 33 | ||
| 34 | typedef enum { | ||
| 35 | DDMT_SSD | ||
| 36 | } DNNDetectionModelType; | ||
| 37 | |||
| 34 | typedef struct DnnDetectContext { | 38 | typedef struct DnnDetectContext { |
| 35 | const AVClass *class; | 39 | const AVClass *class; |
| 36 | DnnContext dnnctx; | 40 | DnnContext dnnctx; |
| @@ -38,6 +42,7 @@ typedef struct DnnDetectContext { · 1 unmodified line | |||
| 38 | char *labels_filename; | 42 | char *labels_filename; |
| 39 | char **labels; | 43 | char **labels; |
| 40 | int label_count; | 44 | int label_count; |
| 45 | DNNDetectionModelType model_type; | ||
| 41 | } DnnDetectContext; | 46 | } DnnDetectContext; |
| 42 | 47 | ||
| 43 | #define OFFSET(x) offsetof(DnnDetectContext, dnnctx.x) | 48 | #define OFFSET(x) offsetof(DnnDetectContext, dnnctx.x) |
| @@ -54,12 +59,14 @@ static const AVOption dnn_detect_options[] = { · 10 unmodified lines | |||
| 54 | DNN_COMMON_OPTIONS | 59 | DNN_COMMON_OPTIONS |
| 55 | { "confidence", "threshold of confidence", OFFSET2(confidence), AV_OPT_TYPE_FLOAT, { .dbl = 0.5 }, 0, 1, FLAGS}, | 60 | { "confidence", "threshold of confidence", OFFSET2(confidence), AV_OPT_TYPE_FLOAT, { .dbl = 0.5 }, 0, 1, FLAGS}, |
| 56 | { "labels", "path to labels file", OFFSET2(labels_filename), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS }, | 61 | { "labels", "path to labels file", OFFSET2(labels_filename), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, FLAGS }, |
| 62 | { "model_type", "DNN detection model type", OFFSET2(model_type), AV_OPT_TYPE_INT, { .i64 = DDMT_SSD }, INT_MIN, INT_MAX, FLAGS, "model_type" }, | ||
| 63 | { "ssd", "output shape [1, 1, N, 7]", 0, AV_OPT_TYPE_CONST, { .i64 = DDMT_SSD }, 0, 0, FLAGS, "model_type" }, | ||
| 57 | { NULL } | 64 | { NULL } |
| 58 | }; | 65 | }; |
| 59 | 66 | ||
| 60 | AVFILTER_DEFINE_CLASS(dnn_detect); | 67 | AVFILTER_DEFINE_CLASS(dnn_detect); |
| 61 | 68 | ||
| 62 | static int dnn_detect_post_proc_ov(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) | 69 | static int dnn_detect_post_proc_ssd(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) |
| 63 | { | 70 | { |
| 64 | DnnDetectContext *ctx = filter_ctx->priv; | 71 | DnnDetectContext *ctx = filter_ctx->priv; |
| 65 | float conf_threshold = ctx->confidence; | 72 | float conf_threshold = ctx->confidence; |
| @@ -67,14 +74,12 @@ static int dnn_detect_post_proc_ov(AVFrame *frame, DNNData *output, AVFilterCont · 1 unmodified line | |||
| 67 | int detect_size = output->width; | 74 | int detect_size = output->width; |
| 68 | float *detections = output->data; | 75 | float *detections = output->data; |
| 69 | int nb_bboxes = 0; | 76 | int nb_bboxes = 0; |
| 70 | AVFrameSideData *sd; | ||
| 71 | AVDetectionBBox *bbox; | ||
| 72 | AVDetectionBBoxHeader *header; | 77 | AVDetectionBBoxHeader *header; |
| 78 | AVDetectionBBox *bbox; | ||
| 73 | 79 | ||
| 74 | sd = av_frame_get_side_data(frame, AV_FRAME_DATA_DETECTION_BBOXES); | 80 | if (output->width != 7) { |
| 75 | if (sd) { | 81 | av_log(filter_ctx, AV_LOG_ERROR, "Model output shape doesn't match ssd requirement.\n"); |
| 76 | av_log(filter_ctx, AV_LOG_ERROR, "already have bounding boxes in side data.\n"); | 82 | return AVERROR(EINVAL); |
| 77 | return -1; | ||
| 78 | } | 83 | } |
| 79 | 84 | ||
| 80 | for (int i = 0; i < proposal_count; ++i) { | 85 | for (int i = 0; i < proposal_count; ++i) { |
| @@ -135,6 +140,29 @@ static int dnn_detect_post_proc_ov(AVFrame *frame, DNNData *output, AVFilterCont · 54 unmodified lines | |||
| 135 | return 0; | 140 | return 0; |
| 136 | } | 141 | } |
| 137 | 142 | ||
| 143 | static int dnn_detect_post_proc_ov(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) | ||
| 144 | { | ||
| 145 | AVFrameSideData *sd; | ||
| 146 | DnnDetectContext *ctx = filter_ctx->priv; | ||
| 147 | int ret = 0; | ||
| 148 | |||
| 149 | sd = av_frame_get_side_data(frame, AV_FRAME_DATA_DETECTION_BBOXES); | ||
| 150 | if (sd) { | ||
| 151 | av_log(filter_ctx, AV_LOG_ERROR, "already have bounding boxes in side data.\n"); | ||
| 152 | return -1; | ||
| 153 | } | ||
| 154 | |||
| 155 | switch (ctx->model_type) { | ||
| 156 | case DDMT_SSD: | ||
| 157 | ret = dnn_detect_post_proc_ssd(frame, output, filter_ctx); | ||
| 158 | if (ret < 0) | ||
| 159 | return ret; | ||
| 160 | break; | ||
| 161 | } | ||
| 162 | |||
| 163 | return 0; | ||
| 164 | } | ||
| 165 | |||
| 138 | static int dnn_detect_post_proc_tf(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) | 166 | static int dnn_detect_post_proc_tf(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) |
| 139 | { | 167 | { |
| 140 | DnnDetectContext *ctx = filter_ctx->priv; | 168 | DnnDetectContext *ctx = filter_ctx->priv; |