resolve merge

Author: misko

data_collection/rpi5_inference/inference_config.yaml

@@ -57,7 +57,7 @@ dry-run: False
 
 data-version: 4
 
-seconds-per-sample: 0.5
+seconds-per-sample: 0.1
 
 
 collector:

spf/data_collector.py

@@ -10,7 +10,9 @@
 from typing import Any, Dict, Optional
 
 import numpy as np
-from attr import dataclass
+#from attr import dataclass
+from dataclasses import dataclass, asdict
+
 from tqdm import tqdm
 
 from spf.dataset.v4_data import v4rx_2xf64_keys, v4rx_f64_keys, v4rx_new_dataset
@@ -216,7 +218,7 @@ def get_rx(self, max_retries=15) -> Dict[str, Any]:
         tries = 0
         while tries < max_retries:
             try:
-                signal_matrix = self.pplus.sdr.rx()
+                signal_matrix = self.pplus.sdr.rx() # complex128 for pluto, eventhough its 12bit TODO
                 rssis = self.pplus.rssis()
                 gains = self.pplus.gains()
                 return {"signal_matrix": signal_matrix, "rssis": rssis, "gains": gains}
@@ -420,7 +422,7 @@ def write_to_record_matrix(self, thread_idx, record_idx, read_thread: ThreadedRX
     def run_inner_collector_thread(self):
         futures = []
         with ThreadPoolExecutorWithQueueSizeLimit(
-            max_workers=6, maxsize=12
+            max_workers=1, maxsize=1
         ) as executor:
             for record_index in tqdm(range(self.yaml_config["n-records-per-receiver"])):
                 for read_thread_idx, read_thread in enumerate(self.read_threads):
@@ -474,12 +476,13 @@ def close(self):
 
 # V4 data format
 class DroneDataCollectorRaw(DataCollector):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, realtime_v5inf, *args, **kwargs):
         super(DroneDataCollectorRaw, self).__init__(
             *args,
             thread_class=ThreadedRXRawV4,
             **kwargs,
         )
+        self.realtime_v5inf=realtime_v5inf
 
     def setup_record_matrix(self):
         if self.data_filename is not None:
@@ -511,12 +514,20 @@ def write_to_record_matrix(self, thread_idx, record_idx, data):
         data.gps_lat = current_pos_heading_and_time["gps"][1]
         data.gps_timestamp = current_pos_heading_and_time["gps_time"]
 
+        if self.realtime_v5inf is not None:
+            print("WRITTING!!")
+            data_dict=asdict(data)
+            data_dict['signal_matrix']=data_dict['signal_matrix'].reshape(1,1,*data_dict['signal_matrix'].shape)
+            print("KEYS IN DICT",data_dict.keys())
+            print("SIG",data_dict['signal_matrix'].shape)
+            self.realtime_v5inf.write_to_idx(record_idx, thread_idx, data_dict)
         if self.data_filename is not None:
             z = self.zarr[f"receivers/r{thread_idx}"]
             z.signal_matrix[record_idx] = data.signal_matrix
             for k in v4rx_f64_keys + v4rx_2xf64_keys:
                 z[k][record_idx] = getattr(data, k)  # getattr(data, k)
 
+
     def close(self):
         self.zarr.store.close()
         self.zarr = None

spf/dataset/segmentation.py

@@ -591,6 +591,17 @@ def segment_session(
         else:
             # If no signal windows were identified, use placeholder values
             segmentation_results["weighted_windows_stats"] = np.array([-1, -1, -1])
+    else:
+        segmentation_results['all_windows_stats']=get_all_windows_stats(v,**kwrgs)[1]
+
+        # Transpose the window statistics for easier processing
+        # all_windows_stats shape is (3, N_windows) where:
+        # - Row 0: Trimmed circular mean of phase difference
+        # - Row 1: Trimmed standard deviation of phase difference
+        # - Row 2: Median absolute signal amplitude
+        segmentation_results["all_windows_stats"] = (
+            segmentation_results["all_windows_stats"].astype(np.float16).T
+        )
 
     # add a singleton in front of each of these
     segmentation_results = {
@@ -600,6 +611,32 @@ def segment_session(
 
     return segmentation_results
 
+def get_all_windows_stats(
+    v,
+    window_size,
+    stride,
+    trim,
+    mean_diff_threshold,
+    max_stddev_threshold,
+    drop_less_than_size,
+    min_abs_signal,
+    steering_vectors=None,  # not used but passed in
+):
+    # Verify the signal matrix has the expected shape (2 antennas)
+    assert v.ndim == 2 and v.shape[0] == 2
+
+    # Calculate phase differences between the two antenna elements
+    # This gives us the phase shift that can be used to determine angle of arrival
+    pd = get_phase_diff(v).astype(np.float32)
+
+    # Calculate statistics for each window:
+    # - Trimmed circular mean of phase differences
+    # - Trimmed standard deviation of phase differences
+    # - Median absolute signal amplitude
+    step_idxs, step_stats = windowed_trimmed_circular_mean_and_stddev(
+        v, pd, window_size=window_size, stride=stride, trim=trim
+    )
+    return step_idxs, step_stats
 
 def simple_segment(
     v,
@@ -658,9 +695,10 @@ def simple_segment(
     # - Trimmed circular mean of phase differences
     # - Trimmed standard deviation of phase differences
     # - Median absolute signal amplitude
-    window_idxs_and_stats = windowed_trimmed_circular_mean_and_stddev(
-        v, pd, window_size=window_size, stride=stride, trim=trim
-    )
+    # window_idxs_and_stats = windowed_trimmed_circular_mean_and_stddev(
+    #     v, pd, window_size=window_size, stride=stride, trim=trim
+    # )
+    window_idxs_and_stats = get_all_windows_stats(v=v,window_size=window_size,stride=stride,trim=trim)
     # window_idxs_and_stats has two components:
     # [0] = list of window indices (start_idx, end_idx)
     # [1] = array of statistics (trimmed_mean, trimmed_stddev, abs_signal_median)

spf/dataset/spf_dataset.py

@@ -456,7 +456,6 @@ def __init__(
         skip_fields: List[
             str
         ] = [],  # Data fields to exclude during loading to save memory
-        n_parallel: int = 20,  # Number of parallel processes for segmentation
         empirical_data_fn: (
             str | None
         ) = None,  # Path to empirical distribution data file for phase-to-angle mapping
@@ -477,7 +476,6 @@ def __init__(
     ):
         # Store configuration parameters
         self.yaml_fn = yaml_fn
-        self.n_parallel = n_parallel
         self.nthetas = nthetas  # Number of angles to discretize space for beamforming
         self.target_ntheta = self.nthetas if target_ntheta is None else target_ntheta
 
@@ -645,14 +643,14 @@ def check_for_new_data(self):
 
     def write_to_idx(self, idx, ridx, raw):
         # this is the heavy lifting of processing, do it on this process
-        rendered_data = self.render_session(idx, ridx, raw)
+        rendered_data = self.render_session(ridx, raw)
 
         self.incoming_queue.put((idx, ridx, rendered_data))
 
         # with self.condition:
         #    self.condition.notify_all()
 
-    def render_session(self, idx, ridx, data):
+    def render_session(self, ridx, data):
         snapshot_idxs = [0]  # which snapshots to get
 
         data["rx_wavelength_spacing"] = torch.tensor(self.rx_wavelength_spacing)
@@ -677,26 +675,29 @@ def render_session(self, idx, ridx, data):
 
         if "signal_matrix" not in self.skip_fields:
             # WARNGING this does not respect flipping!
+            # signal matrix ~ 1,1,2,524288
             abs_signal = data["signal_matrix"].abs().to(torch.float32)
             assert data["signal_matrix"].shape[0] == 1
             pd = torch_get_phase_diff(data["signal_matrix"][0]).to(torch.float32)
             data["abs_signal_and_phase_diff"] = torch.concatenate(
                 [abs_signal, pd[None, :, None]], dim=2
             )
 
-        data["rx_pos_mm"] = torch.vstack(
-            [
-                data["rx_pos_x_mm"],
-                data["rx_pos_y_mm"],
-            ]
-        ).T
+        # data["rx_pos_mm"] = torch.vstack(
+        #     [
+        #         data["rx_pos_x_mm"], # size = [1]
+        #         data["rx_pos_y_mm"], # size = [1]
+        #     ]
+        # ).T # torch.Size([1, 2])
 
-        data["tx_pos_mm"] = torch.vstack(
-            [
-                data["tx_pos_x_mm"],
-                data["tx_pos_y_mm"],
-            ]
-        ).T
+        # data["tx_pos_mm"] = torch.vstack(
+        #     [
+        #         data["tx_pos_x_mm"], # size = [1]
+        #         data["tx_pos_y_mm"], # size = [1]
+        #     ]
+        # ).T # torch.Size([1, 2])
+
+        data["rx_pos_mm"] = data["tx_pos_mm"] = torch.ones(1, 2) * torch.nan
 
         data["rx_pos_xy"] = (
             data["rx_pos_mm"][snapshot_idxs].unsqueeze(0) / self.distance_normalization
@@ -705,9 +706,13 @@ def render_session(self, idx, ridx, data):
         data["tx_pos_xy"] = (
             data["tx_pos_mm"][snapshot_idxs].unsqueeze(0) / self.distance_normalization
         )
-        breakpoint()
+
+        signal_matrix = data["signal_matrix"][0][0]
+        if isinstance(signal_matrix, torch.Tensor):
+            signal_matrix = signal_matrix.numpy()
+
         segmentation = segment_session(
-            data["signal_matrix"][0][0].numpy(),
+            signal_matrix,
             gpu=self.gpu,
             skip_beamformer=False,
             skip_detrend=self.skip_detrend,

spf/mavlink_radio_collection.py

@@ -10,6 +10,8 @@
 from pymavlink import mavutil
 
 from spf.data_collector import DroneDataCollectorRaw
+from spf.dataset.spf_dataset import v5inferencedataset,training_only_keys
+from spf.dataset.spf_nn_dataset_wrapper import v5spfdataset_nn_wrapper
 from spf.distance_finder.distance_finder_controller import DistanceFinderController
 from spf.gps.boundaries import boundaries  # crissy_boundary_convex
 from spf.gps.boundaries import find_closest_boundary
@@ -18,6 +20,7 @@
     drone_get_planner,
     get_ardupilot_serial,
 )
+from spf.scripts.train_utils import load_config_from_fn
 from spf.utils import (
     DataVersionNotImplemented,
     filenames_from_time_in_seconds,
@@ -126,6 +129,23 @@ def parse_args():
         type=str,
         default=None,
     )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default=None,
+    )
+    parser.add_argument(
+        "--checkpoint-config",
+        type=str,
+        default=None,
+    )
+
+    parser.add_argument(
+        "--nthetas",
+        type=int,
+        help="nthetas",
+        default=None,
+    )
     parser.add_argument(
         "--ultrasonic",
         action=argparse.BooleanOptionalAction,
@@ -136,6 +156,11 @@ def parse_args():
         action=argparse.BooleanOptionalAction,
         default=False,
     )
+    parser.add_argument(
+        "--realtime",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
 
     parser.add_argument(
         "--inference", action=argparse.BooleanOptionalAction, default=False
@@ -234,8 +259,42 @@ def parse_args():
     if args.inference:
         pass
 
+
+    if args.checkpoint:
+        # load model config and use that theta
+        config = load_config_from_fn(args.checkpoint_config)
+        assert args.nthetas is None, "nthetas cannot be set when loading checkpoint"
+        args.nthetas = config['global']['nthetas']
+    elif args.nthetas is None:
+        logging.warning("Setting nthetas to 65 as default")
+        args.nthetas=65
+
+    if args.realtime:
+        v5inf = v5inferencedataset(
+            yaml_fn=temp_filenames["yaml"],
+            nthetas=args.nthetas,
+            gpu=False,
+            paired=True,
+            model_config_fn="",
+            skip_fields=["signal_matrix"] + training_only_keys,
+            vehicle_type="rover",
+            skip_segmentation=True,
+            skip_detrend=False,
+        )
+        nn_ds = v5spfdataset_nn_wrapper(
+            v5inf,
+            args.checkpoint_config,
+            args.checkpoint,
+            inference_cache=None,
+            device="cpu",
+            v4=False,
+            absolute=True,
+        )
+
+
     if yaml_config["data-version"] == 4:
         data_collector = DroneDataCollectorRaw(
+            realtime_v5inf=v5inf if args.realtime else None,
             data_filename=temp_filenames["data"] if args.write_to_disk else None,
             yaml_config=yaml_config,
             position_controller=drone,