update: more advanced ML model + support code

Author: fabianschm

e2e.pt

@@ -1,73 +1,58 @@
-import math
-
-import torch
-import numpy as np
-from scipy.signal import butter, filtfilt
-
-import os
-import dotenv
-
-from .csi_database import CSIRecord
-
-dotenv.load_dotenv()
-
-ML_MODEL_SAMPLES_PER_RECORDING = int(os.environ.get("ML_MODEL_SAMPLES_PER_RECORDING"))
-
-
-
-def csi_int16_from_bytes(b: bytes) -> list[int]:
-    return [
-        int.from_bytes(b[2*i:2*i+2], byteorder='little', signed=True)
-        for i in range(len(b) // 2)
-    ]
-
-
-def lowpass_filter(data, cutoff=0.1, order=4):
-    b, a = butter(order, cutoff, btype='low')
-    return filtfilt(b, a, data)
-
-
-def bytes_to_amplitude_phase(raw_csi_bytes: bytes) -> tuple[list[float], list[float]]:
-    csi_fingerprint = csi_int16_from_bytes(raw_csi_bytes)
-    amplitude = []
-    phase = []
-    for j in range(0, 128, 2):
-        real = csi_fingerprint[j]
-        imag = csi_fingerprint[j + 1]
-        amp = math.sqrt(real ** 2 + imag ** 2)
-        phs = math.atan2(imag, real)
-        amplitude.append(amp)  # [64]
-        phase.append(phs)  # [64]
-
-    amplitude = lowpass_filter(amplitude)
-
-    delta_ampl = np.max(amplitude) - np.min(amplitude)
-    amplitude = (amplitude - np.min(amplitude)) / (delta_ampl if delta_ampl != 0 else 1)  # min-max
-
-    phase = list(np.unwrap(phase))
-    delta_phase = np.max(phase) - np.min(phase)
-    phase = (phase - np.min(phase)) / (delta_phase if delta_phase != 0 else 1)
-
-    return amplitude, phase
-
-    #csi_fingerprint = [amplitude, phase]
-
-    #csi_fingerprint = torch.tensor(csi_fingerprint, dtype=torch.float64)  # [time, 2, subcarrier_idx]
-    #print(csi_fingerprint.shape)
-    #csi_fingerprint = csi_fingerprint.transpose(0, 1)  # [2, time, subcarrier_idx]
-    #print(csi_fingerprint.shape)
-
-    #return csi_fingerprint.tolist()
-
-
-def records_to_tensor(csi_records: list[CSIRecord]) -> torch.tensor:
-    amplitudes = []
-    phases = []
-    for record in csi_records:
-        csi_bytes = record.get_csi_bytes()
-        amplitude, phase = bytes_to_amplitude_phase(csi_bytes)
-        amplitudes.append(amplitude)
-        phases.append(phase)
-
-    arr = np.array([[amplitudes, phases]])
-    return torch.tensor(arr, dtype=torch.float)
+import math
+
+import torch
+import numpy as np
+from scipy.signal import butter, filtfilt
+
+import os
+import dotenv
+
+from .csi_database import CSIRecord
+
+dotenv.load_dotenv()
+
+ML_MODEL_SAMPLES_PER_RECORDING = int(os.environ.get("ML_MODEL_SAMPLES_PER_RECORDING"))
+
+
+
+def csi_int16_from_bytes(b: bytes) -> list[int]:
+    return [
+        int.from_bytes(b[2*i:2*i+2], byteorder='little', signed=True)
+        for i in range(len(b) // 2)
+    ]
+
+
+def bytes_to_amplitude_phase(raw_csi_bytes: bytes) -> list[float]:
+    csi_fingerprint = csi_int16_from_bytes(raw_csi_bytes)
+    amplitude = []
+    for j in range(0, 128, 2):
+        if j in [0, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74]:
+            continue
+        real = csi_fingerprint[j]
+        imag = csi_fingerprint[j + 1]
+        amplitude.append(math.sqrt(real ** 2 + imag ** 2))  # [52]
+
+
+    delta_ampl = np.max(amplitude) - np.min(amplitude)
+    amplitude = (amplitude - np.min(amplitude)) / (delta_ampl if delta_ampl != 0 else 1)  # min-max normalization
+
+    return amplitude
+    #csi_fingerprint = [amplitude, phase]
+
+    #csi_fingerprint = torch.tensor(csi_fingerprint, dtype=torch.float64)  # [time, 2, subcarrier_idx]
+    #print(csi_fingerprint.shape)
+    #csi_fingerprint = csi_fingerprint.transpose(0, 1)  # [2, time, subcarrier_idx]
+    #print(csi_fingerprint.shape)
+
+    #return csi_fingerprint.tolist()
+
+
+def records_to_tensor(csi_records: list[CSIRecord]) -> torch.tensor:
+    amplitudes = []
+    for record in csi_records:
+        csi_bytes = record.get_csi_bytes()
+        amplitude = bytes_to_amplitude_phase(csi_bytes)
+        amplitudes.append(amplitude)
+
+    arr = np.array([[amplitudes]])
+    return torch.tensor(arr, dtype=torch.float)

lib/csi_preprocessor.py

@@ -1,83 +1,86 @@
-import torch
-import torch.nn as nn
-import torchvision.models as models
-
-import os
-import dotenv
-
-dotenv.load_dotenv()
-
-ML_MODEL_CHECKPOINT_PATH = os.environ.get("ML_MODEL_CHECKPOINT_PATH")
-
-# taken from https://github.com/RS2002/CrossFi/blob/main/model.py, introduced by "CrossFi: A Cross Domain WiFi Sensing Framework Based on Siamese Network"
-class ResnetBasedSiamese2dNetwork(nn.Module):
-    def __init__(self, output_dims: int = 64, channel: int = 2, pretrained: bool = True, norm: bool = False):
-        super().__init__()
-        self.model = models.resnet18(pretrained)
-        self.model.conv1 = nn.Conv2d(channel, 64, kernel_size=7, stride=2, padding=3, bias=False)
-        self.model.fc = nn.Linear(self.model.fc.in_features, output_dims)
-        self.norm = norm
-
-    def forward(self,x):
-        if self.norm:
-            mean = torch.mean(x, dim=-1, keepdim=True)
-            std = torch.std(x, dim=-1, keepdim=True)
-            y = (x - mean) / std
-        else:
-            y = x
-        return self.model(y)
-
-
-class SiameseDiscriminator(nn.Module):
-    def __init__(self, embedding_dim=64):
-        super().__init__()
-        # after concat you have 2×embedding_dim input
-        self.classifier = nn.Sequential(
-            nn.Linear(2 * embedding_dim, 256),
-            nn.ReLU(inplace=True),
-            nn.Dropout(0.3),
-            nn.Linear(256, 64),
-            nn.ReLU(inplace=True),
-            nn.Linear(64, 1),
-            nn.Sigmoid(),   # outputs P(same)
-        )
-
-    def forward(self, e1, e2):
-        # you could also do torch.abs(e1 - e2) or elementwise mult
-        x = torch.cat([e1, e2], dim=1)
-        return self.classifier(x)
-
-
-class SiameseModel(nn.Module):
-    def __init__(self, embedding_dim=64):
-        super().__init__()
-        self.encoder = ResnetBasedSiamese2dNetwork()
-        self.discriminator = SiameseDiscriminator(embedding_dim)
-
-    def forward(self, x1, x2):
-        e1 = self.encoder(x1)
-        e2 = self.encoder(x2)
-        p_same = self.discriminator(e1, e2)
-        return p_same
-
-
-def model_predict(embedding_model, csi_1, csi_2, device):
-    with torch.no_grad():
-        csi_1, csi_2 = csi_1.to(device), csi_2.to(device)
-
-        p_same = embedding_model(csi_1, csi_2).flatten()
-        pred = (p_same >= 0.9).float().item()
-
-    return pred > 0
-
-
-
-## External Func
-def authenticate(csi_1, csi_2):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    embedding_model = SiameseModel().to(device).eval()
-
-    embedding_model.load_state_dict(torch.load(ML_MODEL_CHECKPOINT_PATH, map_location=device))
-
-    return model_predict(embedding_model, csi_1, csi_2, device)
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+import os
+import dotenv
+
+dotenv.load_dotenv()
+
+ML_MODEL_CHECKPOINT_PATH = os.environ.get("ML_MODEL_CHECKPOINT_PATH")
+
+class CNNEmbedder(nn.Module):
+    def __init__(self, channel: int = 1):
+        super().__init__()
+        self.model = torch.nn.Sequential(
+            nn.Conv2d(channel, 4, kernel_size=3, stride=1), 
+            nn.Tanh(),
+            nn.Dropout2d(p=0.5),
+      
+            nn.Conv2d(4, 8, kernel_size=3, stride=1), 
+            nn.Tanh(),
+            nn.Dropout2d(p=0.5),
+
+            nn.MaxPool2d(kernel_size=3, stride=1),
+
+            nn.Conv2d(8, 16, kernel_size=5, stride=1), 
+            nn.Tanh(),
+            nn.Dropout2d(p=0.5),
+
+            nn.Conv2d(16, 32, kernel_size=3, stride=1),
+            nn.Tanh(),
+            nn.AdaptiveAvgPool2d(1),
+
+            nn.Flatten(),
+        )
+
+    def forward(self,x):
+        return self.model(x)
+
+
+class SiameseDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.classifier = nn.Sequential(
+            nn.Linear(32, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, e1, e2):
+        x = torch.abs(e1 - e2)
+        return self.classifier(x)
+
+
+class SiameseModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = CNNEmbedder()
+        self.discriminator = SiameseDiscriminator()
+
+    def forward(self, x1, x2):
+        e1 = self.encoder(x1)
+        e2 = self.encoder(x2)
+        p_same = self.discriminator(e1, e2)
+        return p_same
+
+
+def model_predict(embedding_model, csi_1, csi_2, device):
+    with torch.no_grad():
+        csi_1, csi_2 = csi_1.to(device), csi_2.to(device)
+
+        p_same = embedding_model(csi_1, csi_2).flatten()
+        pred = (p_same > 0.7).float().item()
+
+    return pred > 0
+
+
+
+## External Func
+def authenticate(csi_1, csi_2):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    embedding_model = SiameseModel().to(device).eval()
+
+    embedding_model.load_state_dict(torch.load(ML_MODEL_CHECKPOINT_PATH, map_location=device))
+
+    return model_predict(embedding_model, csi_1, csi_2, device)