Merge pull request #143 from anarkiwi/diag

handle too small/too low inference result, fix model xywh to opencv rect
IQTLabs · Nov 22, 2023 · ade20ef · ade20ef
2 parents 5355a57 + 96724ad
commit ade20ef
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 21 deletions.
diff --git a/lib/image_inference_impl.cc b/lib/image_inference_impl.cc
@@ -435,37 +435,43 @@ void image_inference_impl::get_inference_() {
           nlohmann::json original_results_json = nlohmann::json::parse(results);
           nlohmann::json results_json = original_results_json;
           size_t rendered_predictions = 0;
-          float xf = float(output_item.points_buffer->cols) /
-                     float(output_item.image_buffer->cols);
-          float yf = float(output_item.points_buffer->rows) /
-                     float(output_item.image_buffer->rows);
+          const float xf = float(output_item.points_buffer->cols) /
+                           float(output_item.image_buffer->cols);
+          const float yf = float(output_item.points_buffer->rows) /
+                           float(output_item.image_buffer->rows);
           const cv::Scalar white = cv::Scalar(255, 255, 255);
           for (auto &prediction_class : original_results_json.items()) {
             size_t i = 0;
             for (auto &prediction_ref : prediction_class.value().items()) {
               auto &prediction = prediction_ref.value();
               float conf = prediction["conf"];
               if (conf > confidence_) {
-                ++rendered_predictions;
                 auto &xywh = prediction["xywh"];
-                int x = xywh[0];
-                int y = xywh[1];
+                int cx = xywh[0];
+                int cy = xywh[1];
                 int w = xywh[2];
                 int h = xywh[3];
-                cv::Mat rssi_points = (*output_item.points_buffer)(cv::Rect(
-                    int(x * xf), int(y * yf), int(w * xf), int(h * yf)));
+                int tlx = cx - (w / 2);
+                int tly = cy - (h / 2);
+                cv::Rect bbox_rect(tlx, tly, w, h);
+                cv::Rect rssi_rect(int(tlx * xf), int(tly * yf), int(w * xf),
+                                   int(h * yf));
+                cv::Mat rssi_points = (*output_item.points_buffer)(rssi_rect);
                 float rssi = cv::mean(rssi_points)[0];
-                auto &augmented = results_json[prediction_class.key()][i];
-                augmented["rssi"] = rssi;
-                augmented["rssi_samples"] = rssi_points.cols * rssi_points.rows;
-                cv::rectangle(*output_item.image_buffer, cv::Rect(x, y, w, h),
-                              white);
-                std::string label = prediction_class.key() + ": conf " +
-                                    std::to_string(conf) + ", RSSI " +
-                                    std::to_string(rssi);
-                cv::putText(*output_item.image_buffer, label,
-                            cv::Point(x - 10, y - 10), cv::FONT_HERSHEY_SIMPLEX,
-                            0.5, white, 2);
+                if (rssi >= min_peak_points_) {
+                  ++rendered_predictions;
+                  auto &augmented = results_json[prediction_class.key()][i];
+                  augmented["rssi"] = rssi;
+                  augmented["rssi_samples"] =
+                      rssi_points.cols * rssi_points.rows;
+                  cv::rectangle(*output_item.image_buffer, bbox_rect, white);
+                  std::string label = prediction_class.key() + ": conf " +
+                                      std::to_string(conf) + ", RSSI " +
+                                      std::to_string(rssi);
+                  cv::putText(*output_item.image_buffer, label,
+                              cv::Point(cx - 10, cy - 10),
+                              cv::FONT_HERSHEY_SIMPLEX, 0.5, white, 2);
+                }
                 // TODO: add NMS
               }
               ++i;

diff --git a/python/iqtlabs/qa_image_inference.py b/python/iqtlabs/qa_image_inference.py
@@ -252,7 +252,7 @@ def predictions_test():
     def test_instance(self):
         port = 11001
         model_name = "testmodel"
-        predictions_result = {"modulation": [{"conf": 0.9, "xywh": [10, 20, 50, 50]}]}
+        predictions_result = {"modulation": [{"conf": 0.9, "xywh": [50, 50, 10, 10]}]}
         if self.pid == 0:
             self.simulate_torchserve(port, model_name, predictions_result)
             return