use common get_samples(), remove multiple parse_filename.

augment/augment.py

@@ -13,62 +13,21 @@
     uniform_discrete_distribution,
 )
 from torchsig.utils.types import SignalData, SignalDescription
-from gamutrf.sample_reader import read_recording, parse_filename
+from gamutrf.sample_reader import read_recording, get_samples
 
 
-def make_signal(samples, sample_rate, center_frequency):
+def make_signal(samples, meta):
     num_iq_samples = samples.shape[0]
     desc = SignalDescription(
-        sample_rate=sample_rate,
+        sample_rate=meta["sample_rate"],
         num_iq_samples=num_iq_samples,
-        center_frequency=center_frequency,
+        center_frequency=meta["center_frequency"],
     )
     # TODO: subclass SignalData with alternate constructor that can take just numpy array
     signal = SignalData(samples.tobytes(), np.float32, np.complex128, desc)
     return signal
 
 
-def get_nosigmf_file(filename):
-    meta = parse_filename(filename)
-    sample_rate = meta["sample_rate"]
-    sample_dtype = meta["sample_dtype"]
-    sample_len = meta["sample_len"]
-    center_frequency = meta["freq_center"]
-    samples = None
-    for samples_buffer in read_recording(
-        filename, sample_rate, sample_dtype, sample_len, max_sample_secs=None
-    ):
-        if samples is None:
-            samples = samples_buffer
-        else:
-            samples = np.concatenate([samples, samples_buffer])
-    signal = make_signal(samples, sample_rate, center_frequency)
-    return filename, signal
-
-
-def get_signal(filename):
-    if not os.path.exists(filename):
-        raise FileNotFoundError(filename)
-    meta_ext = filename.find(".sigmf-meta")
-    if meta_ext == -1:
-        return get_nosigmf_file(filename)
-
-    meta = sigmf.sigmffile.fromfile(filename)
-    data_filename = filename[:meta_ext]
-    meta.set_data_file(data_filename)
-    # read_samples() always converts to host cf32.
-    samples = meta.read_samples()
-    global_meta = meta.get_global_info()
-    sample_rate = global_meta["core:sample_rate"]
-    sample_type = global_meta["core:datatype"]
-    captures_meta = meta.get_captures()
-    center_frequency = None
-    if captures_meta:
-        center_frequency = captures_meta[0].get("core:frequency", None)
-    signal = make_signal(samples, sample_rate, center_frequency)
-    return data_filename, signal
-
-
 def write_signal(filename, signal, transforms_text):
     first_desc = signal.signal_description[0]
     signal.iq_data = signal.iq_data.astype(np.complex64)
@@ -133,7 +92,8 @@ def argument_parser():
 
 def main():
     options = argument_parser().parse_args()
-    data_filename, signal = get_signal(options.filename)
+    data_filename, samples, meta = get_samples(options.filename)
+    signal = make_signal(samples, meta)
     augment(signal, data_filename, options.outdir, options.n, options.transforms)
 
 

gamutrf/grsource.py

@@ -50,7 +50,8 @@ def __init__(
             in_sig=None,
             out_sig=[np.complex64],
         )
-        _, self.samples, self.center_freq = get_samples(input_file)
+        _, self.samples, meta = get_samples(input_file)
+        self.center_freq = meta["center_frequency"]
         self.n_samples = len(self.samples)
         self.i = 0
         self.message_port_register_in(pmt.intern(cmd_port))

gamutrf/offline.py

@@ -7,7 +7,7 @@
 
 from gnuradio import iqtlabs
 from gamutrf.grscan import grscan
-from gamutrf.sample_reader import parse_filename
+from gamutrf.sample_reader import get_samples
 
 
 def argument_parser():
@@ -63,19 +63,20 @@ def main():
     options = argument_parser().parse_args()
     filename = options.filename
     out_dir = os.path.dirname(filename)
-    meta = parse_filename(filename)
+    _data_filename, _samples, meta = get_samples(filename)
+    freq_start = int(meta["center_frequency"] - (meta["sample_rate"] / 2))
     tb = grscan(
         db_clamp_floor=-1e9,
         fft_batch_size=256,
         freq_end=0,
-        freq_start=meta["freq_center"],
+        freq_start=freq_start,
         inference_min_db=-1e9,
         inference_output_dir=out_dir,
         iqtlabs=iqtlabs,
         n_image=options.n_image,
         nfft=options.nfft,
         pretune=True,
-        samp_rate=meta["sample_rate"],
+        samp_rate=int(meta["sample_rate"]),
         sample_dir=out_dir,
         sdr="file:" + filename,
         tune_step_fft=options.tune_step_fft,

gamutrf/sample_reader.py

@@ -8,9 +8,8 @@
 import numpy as np
 from gamutrf.utils import SAMPLE_DTYPES, SAMPLE_FILENAME_RE, is_fft
 
-FFT_FILENAME_RE = re.compile(
-    r"^.+_([0-9]+)_([0-9]+)points_([0-9]+)Hz_([0-9]+)sps\.(s\d+|raw).*$"
-)
+POINTS_RE = re.compile(r"^.+\D([0-9]+)points_.+$")
+
 
 
 def get_reader(filename):
@@ -36,51 +35,41 @@ def default_reader(x):
 
 
 def parse_filename(filename):
+    timestamp = None
+    nfft = None
+    center_frequency = None
+    sample_rate = None
+    sample_type = None
+
     # FFT is always float not matter the original sample type.
     if is_fft(filename):
         sample_type = "raw"
-        match = FFT_FILENAME_RE.match(filename)
-        try:
-            timestamp = int(match.group(1))
-            nfft = int(match.group(2))
-            freq_center = int(match.group(3))
-            sample_rate = int(match.group(4))
-            # sample_type = match.group(3)
-        except AttributeError:
-            timestamp = None
-            nfft = None
-            freq_center = None
-            sample_rate = None
-            sample_type = None
-    else:
-        match = SAMPLE_FILENAME_RE.match(filename)
-        nfft = None
-        try:
-            timestamp = int(match.group(1))
-            freq_center = int(match.group(2))
-            sample_rate = int(match.group(3))
-            sample_type = match.group(4)
-        except AttributeError:
-            timestamp = None
-            freq_center = None
-            sample_rate = None
-            sample_type = None
+        match = POINTS_RE.match(filename)
+        if match is not None:
+            nfft = int(match.group(1))
+
+    match = SAMPLE_FILENAME_RE.match(filename)
+    if match is not None:
+        timestamp, center_frequency, sample_rate, sample_type = (
+            int(match.group(1)),
+            int(match.group(2)),
+            int(match.group(3)),
+            match.group(4),
+        )
 
     sample_dtype, sample_type = SAMPLE_DTYPES.get(sample_type, (None, None))
     sample_bits = None
     sample_len = None
     if sample_dtype:
         if is_fft(filename):
-            sample_dtype = np.float32
-            sample_bits = 32
-            sample_len = 4
+            sample_dtype = np.dtype([("i", np.float32), ("q", np.float32)])
         else:
             sample_dtype = np.dtype([("i", sample_dtype), ("q", sample_dtype)])
-            sample_bits = sample_dtype[0].itemsize * 8
-            sample_len = sample_dtype[0].itemsize * 2
-    file_info = {
+        sample_bits = sample_dtype[0].itemsize * 8
+        sample_len = sample_dtype[0].itemsize * 2
+    meta = {
         "filename": filename,
-        "freq_center": freq_center,
+        "center_frequency": center_frequency,
         "sample_rate": sample_rate,
         "sample_dtype": sample_dtype,
         "sample_len": sample_len,
@@ -89,7 +78,7 @@ def parse_filename(filename):
         "nfft": nfft,
         "timestamp": timestamp,
     }
-    return file_info
+    return meta
 
 
 def read_recording(
@@ -143,19 +132,19 @@ def read_recording(
 
 def get_nosigmf_samples(filename):
     meta = parse_filename(filename)
-    sample_rate = meta["sample_rate"]
-    sample_dtype = meta["sample_dtype"]
-    sample_len = meta["sample_len"]
-    center_frequency = meta["freq_center"]
     samples = None
     for samples_buffer in read_recording(
-        filename, sample_rate, sample_dtype, sample_len, max_sample_secs=None
+        filename,
+        meta["sample_rate"],
+        meta["sample_dtype"],
+        meta["sample_len"],
+        max_sample_secs=None,
     ):
         if samples is None:
             samples = samples_buffer
         else:
             samples = np.concatenate([samples, samples_buffer])
-    return filename, samples, center_frequency
+    return filename, samples, meta
 
 
 def get_samples(filename):
@@ -168,13 +157,18 @@ def get_samples(filename):
     meta = sigmf.sigmffile.fromfile(filename)
     data_filename = filename[:meta_ext]
     meta.set_data_file(data_filename)
-    # read_samples() always converts to host cf32.
     samples = meta.read_samples()
     global_meta = meta.get_global_info()
-    sample_rate = global_meta["core:sample_rate"]
-    sample_type = global_meta["core:datatype"]
     captures_meta = meta.get_captures()
     center_frequency = None
     if captures_meta:
         center_frequency = captures_meta[0].get("core:frequency", None)
-    return data_filename, samples, center_frequency
+    meta = {
+        "sample_rate": global_meta["core:sample_rate"],
+        "sample_dtype": global_meta["core:datatype"],
+        "sample_len": samples[
+            0
+        ].itemsize,  # read_samples() always converts to host cf32.
+        "center_frequency": center_frequency,
+    }
+    return data_filename, samples, meta

gamutrf/samples2raw.py

@@ -1,7 +1,7 @@
 import argparse
 import subprocess
 
-from gamutrf.utils import parse_filename
+from gamutrf.sample_reader import parse_filename
 from gamutrf.utils import replace_ext
 
 
@@ -16,10 +16,11 @@ def make_procs_args(sample_filename, outfmt):
         procs_args.append(["zstdcat", sample_filename])
         out_filename = replace_ext(out_filename, "")
 
-    _, _, sample_rate, _sample_dtype, _sample_len, in_format, in_bits = parse_filename(
-        out_filename
-    )
-    print(_, sample_rate, _sample_dtype, _sample_len, in_format, in_bits)
+    meta = parse_filename(out_filename)
+    sample_rate = meta["sample_rate"]
+    in_format = meta["sample_type"]
+    in_bits = meta["sample_bits"]
+    print(meta)
     out_filename = replace_ext(out_filename, "raw", all_ext=True)
     sox_in = sample_filename
     if procs_args:

gamutrf/specgram.py

@@ -13,9 +13,8 @@
 
 from gamutrf.utils import get_nondot_files
 from gamutrf.utils import is_fft
-from gamutrf.utils import parse_filename
 from gamutrf.utils import replace_ext
-from gamutrf.sample_reader import read_recording
+from gamutrf.sample_reader import read_recording, parse_filename
 
 
 def stride_windows(x, n, noverlap=0, axis=0):
@@ -273,22 +272,16 @@ def plot_spectrogram(
 
 
 def process_recording(args, recording):
-    (
-        _epoch_time,
-        freq_center,
-        sample_rate,
-        sample_dtype,
-        sample_len,
-        _sample_type,
-        _sample_bits,
-    ) = parse_filename(recording)
+    meta = parse_filename(recording)
+    center_frequency = meta["center_frequency"]
+    sample_rate = meta["sample_dtype"]
+    sample_len = meta["sample_len"]
+    sample_dtype = meta["sample_dtype"]
     spectrogram_filename = replace_ext(recording, args.iext, all_ext=True)
     if args.skip_exist and os.path.exists(spectrogram_filename):
         print(f"skipping {recording}")
         return
     if is_fft(recording):
-        # always in complex float format.
-        sample_dtype = np.dtype([("i", "float32"), ("q", "float32")])
         if not args.skip_fft:
             print(f"skipping precomputed FFT {recording}")
             return
@@ -310,7 +303,7 @@ def process_recording(args, recording):
         spectrogram_filename,
         args.nfft,
         sample_rate,
-        freq_center,
+        center_frequency,
         args.cmap,
         args.ytics,
         args.bare,

gamutrf/utils.py

@@ -20,7 +20,7 @@
 # UHD_IMAGES_DIR=/usr/share/uhd/images ./examples/rx_samples_to_file --args num_recv_frames=1000,recv_frame_size=16360 --file test.gz --nsamps 200000000 --rate 20000000 --freq 101e6 --spb 20000000
 ETTUS_ARGS = "num_recv_frames=1000,recv_frame_size=16360,type=b200"
 ETTUS_ANT = "TX/RX"
-SAMPLE_FILENAME_RE = re.compile(r"^.+\D(\d+)_(\d+)Hz_(\d+)sps\.c*([fisu]\d+|raw).*$")
+SAMPLE_FILENAME_RE = re.compile(r"^.+\D(\d+)_(\d+)Hz.*\D(\d+)sps\.c*([fisu]\d+|raw).*$")
 
 SAMPLE_DTYPES = {
     "i8": ("<i1", "signed-integer"),
@@ -83,40 +83,6 @@ def is_fft(filename):
     return os.path.basename(filename).startswith("fft_")
 
 
-def parse_filename(filename):
-    # TODO: parse from sigmf.
-    match = SAMPLE_FILENAME_RE.match(filename)
-    try:
-        epoch_time = int(match.group(1))
-        freq_center = int(match.group(2))
-        sample_rate = int(match.group(3))
-        sample_type = match.group(4)
-    except AttributeError:
-        epoch_time = None
-        freq_center = None
-        sample_rate = None
-        sample_type = None
-    # FFT is always float not matter the original sample type.
-    if is_fft(filename):
-        sample_type = "raw"
-    sample_dtype, sample_type = SAMPLE_DTYPES.get(sample_type, (None, None))
-    sample_bits = None
-    sample_len = None
-    if sample_dtype:
-        sample_dtype = np.dtype([("i", sample_dtype), ("q", sample_dtype)])
-        sample_bits = sample_dtype[0].itemsize * 8
-        sample_len = sample_dtype[0].itemsize * 2
-    return (
-        epoch_time,
-        freq_center,
-        sample_rate,
-        sample_dtype,
-        sample_len,
-        sample_type,
-        sample_bits,
-    )
-
-
 def get_nondot_files(filedir, glob="*.s*.*"):
     return [
         str(path)