Initial torchserve API support.

IQTLabs · Oct 16, 2023 · 46b9cdf · 46b9cdf
1 parent e2230ac
commit 46b9cdf
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 67 deletions.
diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base
@@ -87,7 +87,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
   libvulkan-dev \
   python3-numpy
 WORKDIR /root
-RUN git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.42
+RUN git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.43
 COPY --from=iqtlabs/gamutrf-vkfft:latest /root /root/gr-iqtlabs
 WORKDIR /root/gr-iqtlabs/build
 COPY --from=sigmf-builder /usr/local /usr/local

diff --git a/docs/README-airt.md b/docs/README-airt.md
@@ -85,7 +85,7 @@ install gr-iqtlabs
 $ git clone https://github.com/google/flatbuffers -b v23.5.26
 $ git clone https://github.com/nlohmann/json -b v3.11.2
 $ git clone https://github.com/deepsig/libsigmf -b v1.0.2
-$ git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.42
+$ git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.43
 $ mkdir -p flatbuffers/build && cd flatbuffers/build && cmake -DCMAKE_INSTALL_PREFIX=~/.conda/envs/$CONDA_DEFAULT_ENV .. && make -j $(nproc) && make install && cd ../..
 $ mkdir -p json/build && cd json/build && cmake -DCMAKE_INSTALL_PREFIX=~/.conda/envs/$CONDA_DEFAULT_ENV .. && make -j $(nproc) && make install && cd ../..
 $ mkdir -p libsigmf/build && cd libsigmf/build && cmake -DUSE_SYSTEM_JSON=ON -DUSE_SYSTEM_FLATBUFFERS=ON -DCMAKE_INSTALL_PREFIX=~/.conda/envs/$CONDA_DEFAULT_ENV -DCMAKE_CXX_FLAGS="-I $HOME/.conda/envs/$CONDA_DEFAULT_ENV/include" .. && make -j $(nproc) && make install && cd ../..
@@ -110,7 +110,7 @@ $ pip3 install .
 run gamutrf (may need to change sample rate depending on SDR - e.g. 125e6 or 100e6).
 
 ```
-$ LD_PRELOAD=$HOME/.conda/envs/$CONDA_DEFAULT_ENV/lib/libgomp.so.1 gamutrf-scan --sdr=SoapyAIRT --freq-start=300e6 --freq-end=6e9 --tune-step-fft 256 --samp-rate=100e6 --nfft 256 --pretune
+$ LD_PRELOAD=$HOME/.conda/envs/$CONDA_DEFAULT_ENV/lib/libgomp.so.1 gamutrf-scan --sdr=SoapyAIRT --freq-start=300e6 --freq-end=6e9 --tune-step-fft 256 --samp-rate=100e6 --nfft 256 --pretune --no-tag-now
 ```
 
 gamutrf-scan will repeatedly print

diff --git a/gamutrf/grscan.py b/gamutrf/grscan.py
@@ -20,7 +20,6 @@
     sys.exit(1)
 
 from gamutrf.grsource import get_source
-from gamutrf.gryolo import yolo_bbox
 from gamutrf.utils import endianstr
 
 
@@ -37,8 +36,10 @@ def __init__(
         freq_end=1e9,
         freq_start=100e6,
         igain=0,
+        inference_min_db=-200,
+        inference_model_server="",
+        inference_model_name="",
         inference_output_dir="",
-        inference_plan_file="",
         iqtlabs=None,
         logaddr="0.0.0.0",  # nosec
         logport=8001,
@@ -176,89 +177,43 @@ def __init__(
         self.fft_blocks.append((zeromq.pub_sink(1, 1, zmq_addr, 100, False, 65536, "")))
 
         self.inference_blocks = []
-        if inference_plan_file and inference_output_dir:
-            if not self.wavelearner:
-                raise ValueError(
-                    "trying to use inference but wavelearner not available"
-                )
-            inference_batch_size = 1
+        if inference_output_dir:
             x = 640
             y = 640
             image_shape = (x, y, 3)
             image_vlen = np.prod(image_shape)
             prediction_shape = (1, 8, 8400)
             prediction_vlen = np.prod(prediction_shape)
+            image_dir = Path(inference_output_dir, "images")
             Path(inference_output_dir).mkdir(parents=True, exist_ok=True)
-            Path(inference_output_dir, "images").mkdir(parents=True, exist_ok=True)
+            image_dir.mkdir(parents=True, exist_ok=True)
             self.image_inference_block = self.iqtlabs.image_inference(
                 tag="rx_freq",
                 vlen=nfft,
                 x=x,
                 y=y,
-                image_dir=str(Path(inference_output_dir, "images")),
+                image_dir=str(image_dir),
                 convert_alpha=255,
                 norm_alpha=0,
                 norm_beta=1,
                 norm_type=32,  # cv::NORM_MINMAX = 32
                 colormap=16,  # cv::COLORMAP_VIRIDIS = 16, cv::COLORMAP_TURBO = 20,
                 interpolation=1,  # cv::INTER_LINEAR = 1,
                 flip=0,
+                min_peak_points=inference_min_db,
+                model_server=inference_model_server,
+                model_name=inference_model_name,
             )
-            self.wavelearner_inference_block = self.wavelearner.inference(
-                plan_filepath=inference_plan_file,
-                complex_input=False,
-                input_vlen=inference_batch_size * image_vlen,
-                output_vlen=prediction_vlen * inference_batch_size,
-                batch_size=inference_batch_size,
-            )
-            self.image_to_inference_blocks = [
-                blocks.stream_to_vector(gr.sizeof_char * image_vlen, 1),
-                blocks.vector_to_stream(gr.sizeof_char, image_vlen),
-                blocks.uchar_to_float(),
-                blocks.stream_to_vector(gr.sizeof_float, image_vlen),
-            ]
             self.inference_blocks = [
                 blocks.stream_to_vector(gr.sizeof_float * nfft, 1),
                 self.image_inference_block,
-                *self.image_to_inference_blocks,
-                # FOR DEBUG
-                # iqtlabs.write_freq_samples(
-                #     "rx_freq",
-                #     gr.sizeof_char * image_vlen,
-                #     1,
-                #     inference_output_dir,
-                #     "inference",
-                #     image_vlen,
-                #     0,
-                #     int(samp_rate),
-                # ),
-                self.wavelearner_inference_block,
-                # iqtlabs.write_freq_samples(
-                #     "rx_freq",
-                #     gr.sizeof_float * prediction_vlen * inference_batch_size,
-                #     1,
-                #     inference_output_dir,
-                #     "inference_predictions",
-                #     prediction_vlen * inference_batch_size,
-                #     0,
-                #     int(samp_rate),
-                # ),
+                blocks.file_sink(
+                    gr.sizeof_char,
+                    str(Path(inference_output_dir, "predictions.txt")),
+                    False,
+                ),
             ]
 
-            self.yolo_bbox_block = yolo_bbox(
-                image_shape=image_shape,
-                prediction_shape=prediction_shape,
-                batch_size=inference_batch_size,
-                sample_rate=samp_rate,
-                output_dir=inference_output_dir,
-            )
-            self.connect(
-                (self.image_to_inference_blocks[-1], 0), (self.yolo_bbox_block, 0)
-            )
-            self.connect(
-                (self.wavelearner_inference_block, 0), (self.yolo_bbox_block, 1)
-            )
-
         if pretune:
             self.msg_connect((self.retune_pre_fft, "tune"), (self.sources[0], cmd_port))
             self.msg_connect((self.retune_pre_fft, "tune"), (retune_fft, "cmd"))

diff --git a/gamutrf/gryolo.py b/gamutrf/gryolo.py
@@ -36,6 +36,8 @@ def work(self, input_items, output_items):
         return len(input_items[0])
 
 
+# TODO: refactor to consume predictions feed, and retrieve png file from local disk
+# (provided by image_inference block)
 class yolo_bbox(gr.sync_block):
     def __init__(
         self,

diff --git a/gamutrf/scan.py b/gamutrf/scan.py
@@ -217,11 +217,25 @@ def argument_parser():
         help="Use dc_block_cc long form",
     )
     parser.add_argument(
-        "--inference_plan_file",
-        dest="inference_plan_file",
+        "--inference_min_db",
+        dest="inference_min_db",
+        type=float,
+        default=-50,
+        help="run inference over minimum dB power",
+    )
+    parser.add_argument(
+        "--inference_model_server",
+        dest="inference_model_server",
         type=str,
         default="",
-        help="full path to plan file for wavelearner",
+        help="torchserve model server inference API address (e.g. localhost:1234)",
+    )
+    parser.add_argument(
+        "--inference_model_name",
+        dest="inference_model_name",
+        type=str,
+        default="",
+        help="torchserve model name (e.g. yolov8)",
     )
     parser.add_argument(
         "--inference_output_dir",
@@ -287,7 +301,7 @@ def argument_parser():
     parser.add_argument(
         "--tag-now",
         dest="tag_now",
-        default=False,
+        default=True,
         action=BooleanOptionalAction,
         help="send tag:now command when retuning",
     )

diff --git a/tests/test_grscan.py b/tests/test_grscan.py
@@ -108,6 +108,7 @@ def run_grscan_smoke(self, pretune, wavelearner, write_samples):
                 db_clamp_floor=-1e6,
                 pretune=pretune,
                 fft_batch_size=256,
+                inference_output_dir=str(tempdir),
             )
             tb.start()
             time.sleep(3)