add background rssi stats, add predictions to image, add rssi of pred…

grc/iqtlabs_image_inference.block.yml

@@ -9,7 +9,7 @@ templates:
   make: >
     iqtlabs.image_inference(${tag}, ${vlen}, ${x}, ${y}, ${image_dir},
     ${convert_alpha}, ${norm_alpha}, ${norm_beta}, ${norm_type}, ${colormap},
-    ${interpolation}, ${model_server}, ${model_name})
+    ${interpolation}, ${model_server}, ${model_name}, ${confidence})
 
 cpp_templates:
   includes: ['#include <gnuradio/iqtlabs/image_inference.h>']
@@ -18,7 +18,7 @@ cpp_templates:
     this->${id} = gr::iqtlabs::image_inference::make(${tag}, ${vlen},
     ${x}, ${y}, ${image_dir}, ${convert_alpha}, ${norm_alpha}, ${norm_beta},
     ${norm_type}, ${colormap}, ${interpolation}, ${model_server},
-    ${model_name});
+    ${model_name}, ${confidence});
   link: ['libgnuradio-iqtlabs.so']
 
 
@@ -61,6 +61,8 @@ parameters:
     dtype: str
   - id: model_server
     dtype: str
+  - id: confidence
+    dtype: float
 
 asserts:
   - ${ tag != "" }

include/gnuradio/iqtlabs/image_inference.h

@@ -233,7 +233,7 @@ class IQTLABS_API image_inference : virtual public gr::block {
                    double norm_alpha, double norm_beta, int norm_type,
                    int colormap, int interpolation, int flip,
                    double min_peak_points, const std::string &model_server,
-                   const std::string &model_name);
+                   const std::string &model_name, double confidence);
 };
 
 } // namespace iqtlabs

lib/image_inference_impl.cc

@@ -220,19 +220,19 @@ image_inference::make(const std::string &tag, int vlen, int x, int y,
                       double norm_alpha, double norm_beta, int norm_type,
                       int colormap, int interpolation, int flip,
                       double min_peak_points, const std::string &model_server,
-                      const std::string &model_name) {
+                      const std::string &model_name, double confidence) {
   return gnuradio::make_block_sptr<image_inference_impl>(
       tag, vlen, x, y, image_dir, convert_alpha, norm_alpha, norm_beta,
       norm_type, colormap, interpolation, flip, min_peak_points, model_server,
-      model_name);
+      model_name, confidence);
 }
 
 image_inference_impl::image_inference_impl(
     const std::string &tag, int vlen, int x, int y,
     const std::string &image_dir, double convert_alpha, double norm_alpha,
     double norm_beta, int norm_type, int colormap, int interpolation, int flip,
     double min_peak_points, const std::string &model_server,
-    const std::string &model_name)
+    const std::string &model_name, double confidence)
     : gr::block("image_inference",
                 gr::io_signature::make(1 /* min inputs */, 1 /* max inputs */,
                                        vlen * sizeof(input_type)),
@@ -243,8 +243,9 @@ image_inference_impl::image_inference_impl(
       norm_alpha_(norm_alpha), norm_beta_(norm_beta), norm_type_(norm_type),
       colormap_(colormap), interpolation_(interpolation), flip_(flip),
       min_peak_points_(min_peak_points), model_name_(model_name),
-      running_(true), inference_connected_(false) {
-  points_buffer_ = new cv::Mat(cv::Size(vlen, 0), CV_32F, cv::Scalar::all(0));
+      confidence_(confidence), running_(true), inference_connected_(false) {
+  points_buffer_.reset(
+      new cv::Mat(cv::Size(vlen_, 0), CV_32F, cv::Scalar::all(0)));
   cmapped_buffer_.reset(
       new cv::Mat(cv::Size(vlen, 0), CV_8UC3, cv::Scalar::all(0)));
   resized_buffer_.reset(
@@ -281,7 +282,6 @@ image_inference_impl::~image_inference_impl() {
   while (!inference_q_.empty()) {
     delete_inference_();
   }
-  delete points_buffer_;
 }
 
 void image_inference_impl::process_items_(size_t c, const input_type *&in) {
@@ -298,9 +298,17 @@ void image_inference_impl::create_image_() {
       output_item_type output_item;
       output_item.rx_freq = last_rx_freq_;
       output_item.ts = last_rx_time_;
+      output_item.points_min = points_min;
+      output_item.points_max = points_max;
       output_item.image_buffer =
           new cv::Mat(cv::Size(x_, y_), CV_8UC3, cv::Scalar::all(0));
-      output_item.orig_rows = points_buffer_->rows;
+      output_item.points_buffer =
+          new cv::Mat(cv::Size(vlen_, 0), CV_32F, cv::Scalar::all(0));
+      if (flip_ == -1 || flip_ == 0 || flip_ == 1) {
+        cv::flip(*points_buffer_, *points_buffer_, flip_);
+      }
+      points_buffer_->copyTo(*output_item.points_buffer);
+      output_item.points_mean = cv::mean(*output_item.points_buffer)[0];
       this->d_logger->debug("rx_freq {} rx_time {} rows {}", last_rx_freq_,
                             last_rx_time_, points_buffer_->rows);
       cv::normalize(*points_buffer_, *points_buffer_, norm_alpha_, norm_beta_,
@@ -310,22 +318,37 @@ void image_inference_impl::create_image_() {
       cv::cvtColor(*cmapped_buffer_, *cmapped_buffer_, cv::COLOR_BGR2RGB);
       cv::resize(*cmapped_buffer_, *resized_buffer_, cv::Size(x_, y_),
                  interpolation_);
-      if (flip_ == -1 || flip_ == 0 || flip_ == 1) {
-        cv::flip(*resized_buffer_, *resized_buffer_, flip_);
-      }
       cv::cvtColor(*resized_buffer_, *output_item.image_buffer,
                    cv::COLOR_RGB2BGR);
-      output_item.points_buffer = points_buffer_;
       if (!inference_q_.push(output_item)) {
         d_logger->error("inference request queue full, size {}", MAX_INFERENCE);
         delete_output_item_(output_item);
       }
     }
-    points_buffer_ =
-        new cv::Mat(cv::Size(vlen_, 0), CV_32F, cv::Scalar::all(0));
+    points_buffer_->resize(0);
   }
 }
 
+std::string image_inference_impl::write_image_(
+    const std::string &prefix, output_item_type &output_item,
+    boost::scoped_ptr<std::vector<unsigned char>> &encoded_buffer) {
+  encoded_buffer.reset(new std::vector<unsigned char>());
+  cv::imencode(IMAGE_EXT, *output_item.image_buffer, *encoded_buffer);
+  std::string image_file_base =
+      prefix + "_" + host_now_str_(output_item.ts) + "_" +
+      std::to_string(uint64_t(x_)) + "x" + std::to_string(uint64_t(y_)) + "_" +
+      std::to_string(uint64_t(output_item.rx_freq)) + "Hz";
+  std::string image_file_png = image_file_base + IMAGE_EXT;
+  std::string dot_image_file_png = image_dir_ + "/." + image_file_png;
+  std::string full_image_file_png = image_dir_ + "/" + image_file_png;
+  std::ofstream image_out;
+  image_out.open(dot_image_file_png, std::ios::binary | std::ios::out);
+  image_out.write((const char *)encoded_buffer->data(), encoded_buffer->size());
+  image_out.close();
+  rename(dot_image_file_png.c_str(), full_image_file_png.c_str());
+  return full_image_file_png;
+}
+
 void image_inference_impl::get_inference_() {
   boost::beast::error_code ec;
   for (;;) {
@@ -342,30 +365,20 @@ void image_inference_impl::get_inference_() {
     }
     output_item_type output_item;
     inference_q_.pop(output_item);
-    boost::scoped_ptr<std::vector<unsigned char>> encoded_buffer(
-        new std::vector<unsigned char>());
-    cv::imencode(IMAGE_EXT, *output_item.image_buffer, *encoded_buffer);
-    // write image file
-    std::string image_file_base =
-        "image_" + host_now_str_(output_item.ts) + "_" +
-        std::to_string(uint64_t(x_)) + "x" + std::to_string(uint64_t(y_)) +
-        "_" + std::to_string(uint64_t(output_item.rx_freq)) + "Hz";
-    std::string image_file_png = image_file_base + IMAGE_EXT;
-    std::string dot_image_file_png = image_dir_ + "/." + image_file_png;
-    std::string full_image_file_png = image_dir_ + "/" + image_file_png;
-    std::ofstream image_out;
-    image_out.open(dot_image_file_png, std::ios::binary | std::ios::out);
-    image_out.write((const char *)encoded_buffer->data(),
-                    encoded_buffer->size());
-    image_out.close();
-    rename(dot_image_file_png.c_str(), full_image_file_png.c_str());
+    boost::scoped_ptr<std::vector<unsigned char>> encoded_buffer;
+
+    nlohmann::json metadata_json;
+    metadata_json["rssi_max"] = std::to_string(output_item.points_max);
+    metadata_json["rssi_mean"] = std::to_string(output_item.points_mean);
+    metadata_json["rssi_min"] = std::to_string(output_item.points_min);
+    metadata_json["ts"] = host_now_str_(output_item.ts);
+    metadata_json["rx_freq"] = std::to_string(output_item.rx_freq);
+    metadata_json["orig_rows"] = output_item.points_buffer->rows;
+    metadata_json["image_path"] =
+        write_image_("image", output_item, encoded_buffer);
+
+    nlohmann::json output_json;
 
-    std::stringstream ss("", std::ios_base::app | std::ios_base::out);
-    ss << "{"
-       << "\"ts\": " << host_now_str_(output_item.ts)
-       << ", \"rx_freq\": " << output_item.rx_freq
-       << ", \"orig_rows\": " << output_item.orig_rows << ", \"image_path\": \""
-       << full_image_file_png << "\"";
     if (host_.size() && port_.size()) {
       const std::string_view body(
           reinterpret_cast<char const *>(encoded_buffer->data()),
@@ -412,35 +425,67 @@ void image_inference_impl::get_inference_() {
           boost::beast::http::read(*stream_, buffer, res);
           results = res.body().data();
         } catch (std::exception &ex) {
-          ss << ", \"error\": \"" << ex.what() << "\"";
+          output_json["error"] = ex.what();
           inference_connected_ = false;
         }
       }
 
       if (results.size()) {
         if (nlohmann::json::accept(results)) {
-          nlohmann::json results_json = nlohmann::json::parse(results);
-          ss << ", \"predictions\": " << results_json;
-          // TODO: create with predictions and bboxes.
-          // TODO: extract RSSI via bounding box
-          // cv::Rect rect(100, 100, 50, 50);
-          // cv::rectangle(*resized_buffer_, rect, cv::Scalar(255, 255, 255));
-          // cv::imencode(IMAGE_EXT, *resized_buffer_,
-          // *output_item.image_buffer); image_out.open("/tmp/test.png",
-          // std::ios::binary | std::ios::out); image_out.write((const
-          // char*)output_item.image_buffer->data(),
-          //                 output_item.image_buffer->size());
-          // image_out.close();
+          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 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 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)));
+                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);
+                // TODO: add NMS
+              }
+              ++i;
+            }
+          }
+          output_json["predictions"] = results_json;
+          if (rendered_predictions) {
+            metadata_json["predictions_image_path"] =
+                write_image_("predictions_image", output_item, encoded_buffer);
+          }
         } else {
-          ss << ", \"error\": \"invalid json: " << results << "\"";
+          output_json["error"] = "invalid json: " + results;
           inference_connected_ = false;
         }
       }
     }
-    // double new line to faciliate json parsing, since prediction may contain
+    // double new line to facilitate json parsing, since prediction may contain
     // new lines.
-    ss << "}\n" << std::endl;
-    json_q_.push(ss.str());
+    output_json["metadata"] = metadata_json;
+    json_q_.push(output_json.dump() + "\n\n");
     delete_output_item_(output_item);
   }
 }

lib/image_inference_impl.h

@@ -232,19 +232,20 @@ typedef struct output_item {
   double ts;
   cv::Mat *image_buffer;
   cv::Mat *points_buffer;
-  size_t orig_rows;
+  double points_min;
+  double points_mean;
+  double points_max;
 } output_item_type;
 
 class image_inference_impl : public image_inference, base_impl {
 private:
   int x_, y_, vlen_, norm_type_, colormap_, interpolation_, flip_;
   uint64_t last_rx_freq_;
   double convert_alpha_, norm_alpha_, norm_beta_, last_rx_time_,
-      min_peak_points_;
+      min_peak_points_, confidence_;
   boost::lockfree::spsc_queue<output_item_type> inference_q_{MAX_INFERENCE};
   boost::lockfree::spsc_queue<std::string> json_q_{MAX_INFERENCE};
-  cv::Mat *points_buffer_;
-  boost::scoped_ptr<cv::Mat> cmapped_buffer_, resized_buffer_;
+  boost::scoped_ptr<cv::Mat> cmapped_buffer_, resized_buffer_, points_buffer_;
   std::string image_dir_;
   pmt::pmt_t tag_;
   std::deque<output_type> out_buf_;
@@ -260,14 +261,17 @@ class image_inference_impl : public image_inference, base_impl {
   void get_inference_();
   void delete_output_item_(output_item_type &output_item);
   void delete_inference_();
+  std::string
+  write_image_(const std::string &prefix, output_item_type &output_item,
+               boost::scoped_ptr<std::vector<unsigned char>> &encoded_buffer);
 
 public:
   image_inference_impl(const std::string &tag, int vlen, int x, int y,
                        const std::string &image_dir, double convert_alpha,
                        double norm_alpha, double norm_beta, int norm_type,
                        int colormap, int interpolation, int flip,
                        double min_peak_points, const std::string &model_server,
-                       const std::string &model_name);
+                       const std::string &model_name, double confidence);
   ~image_inference_impl();
   int general_work(int noutput_items, gr_vector_int &ninput_items,
                    gr_vector_const_void_star &input_items,

python/iqtlabs/bindings/image_inference_python.cc

@@ -14,7 +14,7 @@
 /* BINDTOOL_GEN_AUTOMATIC(0)                                                       */
 /* BINDTOOL_USE_PYGCCXML(0)                                                        */
 /* BINDTOOL_HEADER_FILE(image_inference.h)                                        */
-/* BINDTOOL_HEADER_FILE_HASH(ed9a9d38e8c7f50546d907bfa44a223b)                     */
+/* BINDTOOL_HEADER_FILE_HASH(95fb94fd34e1930b2cdc80357e966443)                     */
 /***********************************************************************************/
 
 #include <pybind11/complex.h>
@@ -54,6 +54,7 @@ void bind_image_inference(py::module& m)
              py::arg("min_peak_points"),
              py::arg("model_server"),
              py::arg("model_name"),
+             py::arg("confidence"),
              D(image_inference, make))
 
 

python/iqtlabs/qa_image_inference.py

@@ -252,15 +252,15 @@ def predictions_test():
     def test_instance(self):
         port = 11001
         model_name = "testmodel"
-        predictions_result = {"modulation": [{"conf": 0.9, "xywh": [1, 2, 3, 4]}]}
+        predictions_result = {"modulation": [{"conf": 0.9, "xywh": [10, 20, 50, 50]}]}
         if self.pid == 0:
             self.simulate_torchserve(port, model_name, predictions_result)
             return
         x = 800
         y = 600
         fft_size = 1024
         output_vlen = x * y * 3
-        samp_rate = 32e3
+        samp_rate = 4e6
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = os.path.join(tmpdir, "samples")
             freq_divisor = 1e9
@@ -285,6 +285,7 @@ def test_instance(self):
                 -1e9,
                 f"localhost:{port}",
                 model_name,
+                0.8,
             )
             c2r = blocks.complex_to_real(1)
             stream2vector = blocks.stream_to_vector(gr.sizeof_float, fft_size)
@@ -304,7 +305,7 @@ def test_instance(self):
             time.sleep(test_time)
             self.tb.stop()
             self.tb.wait()
-            image_files = [f for f in glob.glob(f"{tmpdir}/image*png")]
+            image_files = [f for f in glob.glob(f"{tmpdir}/*image*png")]
             self.assertGreater(len(image_files), 2)
             for image_file in image_files:
                 stat = os.stat(image_file)
@@ -319,7 +320,18 @@ def test_instance(self):
                 if not json_raw:
                     continue
                 result = json.loads(json_raw)
-                self.assertTrue(os.path.exists(result["image_path"]))
+                print(result)
+                metadata_result = result["metadata"]
+                rssi_min, rssi_mean, rssi_max = [
+                    float(metadata_result[v])
+                    for v in ("rssi_min", "rssi_mean", "rssi_max")
+                ]
+                self.assertGreaterEqual(rssi_mean, rssi_min, metadata_result)
+                self.assertGreaterEqual(rssi_max, rssi_mean, metadata_result)
+                self.assertTrue(os.path.exists(metadata_result["image_path"]))
+                self.assertTrue(os.path.exists(metadata_result["predictions_image_path"]))
+                for k in ("rssi", "rssi_samples"):
+                    del result["predictions"]["modulation"][0][k]
                 self.assertEqual(result["predictions"], predictions_result)