From 8a99f497f50afbf319f43cd26d0a329dc5975b4f Mon Sep 17 00:00:00 2001
From: Lucas Tindall <ltindall@ucsd.edu>
Date: Wed, 20 Dec 2023 11:23:30 -0800
Subject: [PATCH] Expose colormap option in scanner

---
 gamutrf/grscan.py | 3 ++-
 gamutrf/scan.py   | 8 ++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/gamutrf/grscan.py b/gamutrf/grscan.py
index 1d6e3597..baa6799c 100644
--- a/gamutrf/grscan.py
+++ b/gamutrf/grscan.py
@@ -27,6 +27,7 @@ class grscan(gr.top_block):
     def __init__(
         self,
         bucket_range=1.0,
+        colormap=16,
         compass=False,
         db_clamp_ceil=50,
         db_clamp_floor=-200,
@@ -217,7 +218,7 @@ def __init__(
                         norm_alpha=0,
                         norm_beta=1,
                         norm_type=32,  # cv::NORM_MINMAX = 32
-                        colormap=16,  # cv::COLORMAP_VIRIDIS = 16, cv::COLORMAP_TURBO = 20,
+                        colormap=colormap,  # cv::COLORMAP_VIRIDIS = 16, cv::COLORMAP_TURBO = 20,
                         interpolation=1,  # cv::INTER_LINEAR = 1,
                         flip=0,  # 0 means flipping around the x-axis
                         min_peak_points=inference_min_db,
diff --git a/gamutrf/scan.py b/gamutrf/scan.py
index 7945379d..7df8558e 100644
--- a/gamutrf/scan.py
+++ b/gamutrf/scan.py
@@ -380,6 +380,14 @@ def argument_parser():
         default=0,
         help="if > 1, only send 1/n images for inference",
     )
+    parser.add_argument(
+        "--colormap",
+        dest="colormap",
+        type=int,
+        default=16,
+        help="""OpenCV colormap for spectrograms. See
+        https://docs.opencv.org/4.x/d3/d50/group__imgproc__colormap.html""",
+    )
     return parser