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 | 31 | #include "libavutil/avstring.h" |
| 32 | 32 | #include "libavutil/detection_bbox.h" |
| 33 | 33 | |
| 34 | +typedef enum { | |
| 35 | + DDMT_SSD | |
| 36 | +} DNNDetectionModelType; | |
| 37 | + | |
| 34 | 38 | typedef struct DnnDetectContext { |
| 35 | 39 | const AVClass *class; |
| 36 | 40 | DnnContext dnnctx; |
| @@ -38,6 +42,7 @@ typedef struct DnnDetectContext { · 1 unmodified line | ||
| 38 | 42 | char *labels_filename; |
| 39 | 43 | char **labels; |
| 40 | 44 | int label_count; |
| 45 | + DNNDetectionModelType model_type; | |
| 41 | 46 | } DnnDetectContext; |
| 42 | 47 | |
| 43 | 48 | #define OFFSET(x) offsetof(DnnDetectContext, dnnctx.x) |
| @@ -54,12 +59,14 @@ static const AVOption dnn_detect_options[] = { · 10 unmodified lines | ||
| 54 | 59 | DNN_COMMON_OPTIONS |
| 55 | 60 | { "confidence", "threshold of confidence", OFFSET2(confidence), AV_OPT_TYPE_FLOAT, { .dbl = 0.5 }, 0, 1, FLAGS}, |
| 56 | 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 | 64 | { NULL } |
| 58 | 65 | }; |
| 59 | 66 | |
| 60 | 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 | 71 | DnnDetectContext *ctx = filter_ctx->priv; |
| 65 | 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 | 74 | int detect_size = output->width; |
| 68 | 75 | float *detections = output->data; |
| 69 | 76 | int nb_bboxes = 0; |
| 70 | − AVFrameSideData *sd; | |
| 71 | − AVDetectionBBox *bbox; | |
| 72 | 77 | AVDetectionBBoxHeader *header; |
| 78 | + AVDetectionBBox *bbox; | |
| 73 | 79 | |
| 74 | − sd = av_frame_get_side_data(frame, AV_FRAME_DATA_DETECTION_BBOXES); | |
| 75 | − if (sd) { | |
| 76 | − av_log(filter_ctx, AV_LOG_ERROR, "already have bounding boxes in side data.\n"); | |
| 77 | − return -1; | |
| 80 | + if (output->width != 7) { | |
| 81 | + av_log(filter_ctx, AV_LOG_ERROR, "Model output shape doesn't match ssd requirement.\n"); | |
| 82 | + return AVERROR(EINVAL); | |
| 78 | 83 | } |
| 79 | 84 | |
| 80 | 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 | 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 | 166 | static int dnn_detect_post_proc_tf(AVFrame *frame, DNNData *output, AVFilterContext *filter_ctx) |
| 139 | 167 | { |
| 140 | 168 | DnnDetectContext *ctx = filter_ctx->priv; |