Use the correct size for predictions

RidgeRun · Feb 6, 2020 · 1b20c3a · 1b20c3a
1 parent b22ac76
commit 1b20c3a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/gst-libs/gst/r2inference/gstbackend.cc b/gst-libs/gst/r2inference/gstbackend.cc
@@ -432,8 +432,8 @@ gst_backend_process_frame (GstBackend *self, GstVideoFrame *input_frame,
   *prediction_size = prediction->GetResultSize ();
 
   /*could we avoid memory copy ?*/
-  *prediction_data = g_malloc(*prediction_size);
-  memcpy(*prediction_data, prediction->GetResultData(), *prediction_size);
+  *prediction_data = g_malloc(*prediction_size * sizeof(gfloat));
+  memcpy(*prediction_data, prediction->GetResultData(), *prediction_size * sizeof(gfloat));
 
   GST_LOG_OBJECT (self, "Size of prediction %p is %lu",
                   *prediction_data, *prediction_size);

diff --git a/gst-libs/gst/r2inference/gstinferencepostprocess.c b/gst-libs/gst/r2inference/gstinferencepostprocess.c
@@ -48,7 +48,7 @@ gst_fill_classification_meta (GstClassificationMeta * class_meta,
   g_return_val_if_fail (class_meta != NULL, FALSE);
   g_return_val_if_fail (prediction != NULL, FALSE);
 
-  class_meta->num_labels = predsize / sizeof (gfloat);
+  class_meta->num_labels = predsize;
   class_meta->label_probs =
       g_malloc (class_meta->num_labels * sizeof (gdouble));
   for (gint i = 0; i < class_meta->num_labels; ++i) {