CMuSeNet Training / Validation code and Synthetic IQ samples generator

CMuSeNet_BIGRED.ipynb

@@ -0,0 +1,1277 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b5007b71",
+   "metadata": {},
+   "source": [
+    "### Initialization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e6b1226",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "from scipy.signal import welch\n",
+    "import torch\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from tqdm import tqdm\n",
+    "import math\n",
+    "import json\n",
+    "\n",
+    "# Constants\n",
+    "START_INDEX = 10  # Skip first few samples\n",
+    "SIGNAL_LENGTH = 1024 * 16\n",
+    "SAMPLE_RATE = 20e6\n",
+    "MASK_SIZE = 1024 * 16  # Mask size for segmentation\n",
+    "\n",
+    "# Functions for Signal Processing\n",
+    "def load_real_data(sample_path):\n",
+    "    \"\"\"\n",
+    "    Load raw signal data from a .dat file.\n",
+    "    \"\"\"\n",
+    "    with open(sample_path, \"rb\") as f:\n",
+    "        signal = np.fromfile(f, dtype=np.complex64)\n",
+    "    return signal\n",
+    "\n",
+    "def load_data(signal_id):\n",
+    "    \"\"\"\n",
+    "    Load signal data and its corresponding metadata.\n",
+    "    \"\"\"\n",
+    "    signal = load_real_data(signal_id)\n",
+    "    metadata_file = signal_id.with_suffix(\".json\")\n",
+    "    if metadata_file.exists():\n",
+    "        with open(metadata_file, \"r\") as f:\n",
+    "            metadata = json.load(f)\n",
+    "    else:\n",
+    "        raise FileNotFoundError(f\"Metadata file {metadata_file} not found for signal {signal_id}\")\n",
+    "    return signal[START_INDEX:], metadata, metadata_file\n",
+    "\n",
+    "def apply_psd(signal, Fs, NFFT):\n",
+    "    \"\"\"\n",
+    "    Calculate the PSD and corresponding frequencies using Welch's method.\n",
+    "    \"\"\"\n",
+    "    freqs, psd = welch(signal, fs=Fs, nfft=NFFT, return_onesided=False)\n",
+    "    psd = np.fft.fftshift(psd)\n",
+    "    freqs = np.fft.fftshift(freqs)\n",
+    "    return psd, freqs\n",
+    "\n",
+    "def calculate_fft(signal):\n",
+    "    \"\"\"\n",
+    "    Calculate the FFT of the signal and return real and imaginary parts as separate channels.\n",
+    "    \"\"\"\n",
+    "    signal = signal[:SIGNAL_LENGTH]\n",
+    "    signal = np.fft.fft(signal)\n",
+    "    signal = np.fft.fftshift(signal)\n",
+    "    signal /= np.max(np.abs(signal))\n",
+    "    return signal"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "440b802c",
+   "metadata": {},
+   "source": [
+    "### Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31bc3770",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Dataset Class\n",
+    "class WidebandSignalDataset(Dataset):\n",
+    "    def __init__(self, signal_ids, mask_size=1024 * 16):\n",
+    "        \"\"\"\n",
+    "        Initialize the dataset with signal IDs and the specified mask size.\n",
+    "        \"\"\"\n",
+    "        self.mask_size = mask_size\n",
+    "        self.signal_ids = signal_ids\n",
+    "        self.loaded_data = [self.process_signal(signal_id) for signal_id in tqdm(self.signal_ids)]\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.signal_ids)\n",
+    "\n",
+    "    def __getitem__(self, index):\n",
+    "        return self.loaded_data[index]\n",
+    "\n",
+    "    def process_signal(self, signal_id):\n",
+    "        signal, metadata, _ = load_data(signal_id)\n",
+    "\n",
+    "        # Ensure signal length matches SIGNAL_LENGTH\n",
+    "        if len(signal) < SIGNAL_LENGTH:\n",
+    "            # Pad with zeros if the signal is shorter\n",
+    "            signal = np.pad(signal, (0, SIGNAL_LENGTH - len(signal)), mode='constant')\n",
+    "        elif len(signal) > SIGNAL_LENGTH:\n",
+    "            # Truncate if the signal is longer\n",
+    "            signal = signal[:SIGNAL_LENGTH]\n",
+    "\n",
+    "        # Apply FFT\n",
+    "        signal = np.fft.fft(signal)\n",
+    "        signal = np.fft.fftshift(signal)\n",
+    "        signal /= np.max(np.abs(signal))  # Normalize\n",
+    "        complex_signal = torch.from_numpy(signal).type(torch.complex64).unsqueeze(0)  # Add channel dimension\n",
+    "\n",
+    "        # Create mask with fixed size\n",
+    "        masks = torch.zeros(self.mask_size, dtype=torch.float32)\n",
+    "        scale_ratio = self.mask_size / SAMPLE_RATE\n",
+    "        scaled_metadata = process_metadata(metadata)\n",
+    "        for meta in scaled_metadata:\n",
+    "            f1, f2 = meta[\"position\"]\n",
+    "            x1 = int(math.floor(f1 * scale_ratio))\n",
+    "            x2 = int(math.ceil(f2 * scale_ratio))\n",
+    "            masks[x1:x2] = 1\n",
+    "\n",
+    "        return complex_signal, masks\n",
+    "\n",
+    "\n",
+    "\n",
+    "def process_metadata(metadata):\n",
+    "    \"\"\"\n",
+    "    Scale metadata to the dataset's frequency and bandwidth ranges.\n",
+    "    \"\"\"\n",
+    "    scaled_metadata = [\n",
+    "        {\n",
+    "            \"position\": (\n",
+    "                math.floor((SAMPLE_RATE / 2 + i[\"fc\"] - i[\"bw\"] / 2) * SIGNAL_LENGTH / SAMPLE_RATE),\n",
+    "                math.ceil((SAMPLE_RATE / 2 + i[\"fc\"] + i[\"bw\"] / 2) * SIGNAL_LENGTH / SAMPLE_RATE)\n",
+    "            ),\n",
+    "            \"snr\": 1,  # Placeholder value\n",
+    "            \"bw\": i[\"bw\"],\n",
+    "            \"num\": len(metadata),\n",
+    "            \"esn0\": 1,  # Placeholder value\n",
+    "        }\n",
+    "        for i in metadata\n",
+    "    ]\n",
+    "    return scaled_metadata\n",
+    "\n",
+    "# Dataset Splitting and Initialization\n",
+    "NEW_DATA_DIR = Path(\"/data/bigred/ofh/0\")\n",
+    "def get_real_signals(freq_directory):\n",
+    "    return list(freq_directory.rglob(\"*.dat\"))\n",
+    "\n",
+    "signal_dirs = get_real_signals(NEW_DATA_DIR)\n",
+    "total_signals = len(signal_dirs)\n",
+    "\n",
+    "train_split = int(0.80 * total_signals)\n",
+    "validation_split = int(0.90 * total_signals)\n",
+    "\n",
+    "train, validation, test = (\n",
+    "    signal_dirs[:train_split],\n",
+    "    signal_dirs[train_split:validation_split],\n",
+    "    signal_dirs[validation_split:]\n",
+    ")\n",
+    "\n",
+    "print(f\"Train set size: {len(train)}\")\n",
+    "print(f\"Validation set size: {len(validation)}\")\n",
+    "print(f\"Test set size: {len(test)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f5305642",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Data Loaders\n",
+    "BATCH_SIZE = 64\n",
+    "\n",
+    "train_dataset = WidebandSignalDataset(signal_ids=train)\n",
+    "validation_dataset = WidebandSignalDataset(signal_ids=validation)\n",
+    "test_dataset = WidebandSignalDataset(signal_ids=test)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54a4f325",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
+    "valid_loader = DataLoader(validation_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
+    "test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3893c583",
+   "metadata": {},
+   "source": [
+    "### CV-ResNet-18"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bc2001c4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import complexPyTorch.complexLayers as cplx\n",
+    "from typing import Optional, Callable, Type, Union, List\n",
+    "import torch.nn.functional as F\n",
+    "from torch import Tensor\n",
+    "\n",
+    "def conv3x3(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"3x3 convolution with padding\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
+    "\n",
+    "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"1x1 convolution\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
+    "\n",
+    "class BasicBlock(nn.Module):\n",
+    "    expansion = 1\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(BasicBlock, self).__init__()\n",
+    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.conv2 = conv3x3(planes, planes)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class Bottleneck(nn.Module):\n",
+    "    expansion = 4\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(Bottleneck, self).__init__()\n",
+    "        self.conv1 = conv1x1(inplanes, planes)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv2 = conv3x3(planes, planes, stride)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv3 = conv1x1(planes, planes * self.expansion)\n",
+    "        self.bn3 = cplx.ComplexBatchNorm2d(planes * self.expansion)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv3(out)\n",
+    "        out = self.bn3(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class ComplexResNet(nn.Module):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        block: Type[Union[BasicBlock, Bottleneck]],\n",
+    "        layers: List[int],\n",
+    "        num_classes: int = SIGNAL_LENGTH,\n",
+    "        zero_init_residual: bool = False,\n",
+    "        groups: int = 1,\n",
+    "        width_per_group: int = 64,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(ComplexResNet, self).__init__()\n",
+    "        if norm_layer is None:\n",
+    "            norm_layer = cplx.ComplexBatchNorm2d\n",
+    "        self._norm_layer = norm_layer\n",
+    "\n",
+    "        self.inplanes = 64\n",
+    "        self.dilation = 1\n",
+    "\n",
+    "        self.groups = groups\n",
+    "        self.base_width = width_per_group\n",
+    "        self.conv1 = cplx.ComplexConv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)\n",
+    "        self.bn1 = norm_layer(self.inplanes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
+    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
+    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
+    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
+    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
+    "        self.fc = cplx.ComplexLinear(512 * block.expansion, num_classes)\n",
+    "        self.sigmoid = cplx.ComplexSigmoid()\n",
+    "\n",
+    "    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int, stride: int = 1) -> nn.Sequential:\n",
+    "        norm_layer = self._norm_layer\n",
+    "        downsample = None\n",
+    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
+    "            downsample = nn.Sequential(\n",
+    "                conv1x1(self.inplanes, planes * block.expansion, stride),\n",
+    "                norm_layer(planes * block.expansion),\n",
+    "            )\n",
+    "\n",
+    "        layers = []\n",
+    "        layers.append(block(self.inplanes, planes, stride, downsample, norm_layer))\n",
+    "        self.inplanes = planes * block.expansion\n",
+    "        for _ in range(1, blocks):\n",
+    "            layers.append(block(self.inplanes, planes, norm_layer=norm_layer))\n",
+    "\n",
+    "        return nn.Sequential(*layers)\n",
+    "\n",
+    "    def _forward_impl(self, x: Tensor) -> Tensor:\n",
+    "        x = self.conv1(x)\n",
+    "        x = self.bn1(x)\n",
+    "        x = self.relu(x)\n",
+    "\n",
+    "        x = self.layer1(x)\n",
+    "        x = self.layer2(x)\n",
+    "        x = self.layer3(x)\n",
+    "        x = self.layer4(x)\n",
+    "\n",
+    "        x = self.avgpool(x)\n",
+    "        x = torch.flatten(x, 1)\n",
+    "        x = self.fc(x)\n",
+    "        x = self.sigmoid(x)\n",
+    "        return x\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        return self._forward_impl(x)\n",
+    "\n",
+    "def ComplexResNet18():\n",
+    "    return ComplexResNet(BasicBlock, [2, 2, 2, 2])\n",
+    "\n",
+    "# Create the model instance\n",
+    "model = ComplexResNet18()\n",
+    "print(model)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a8e09e4",
+   "metadata": {},
+   "source": [
+    "### Early Stop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24f79a24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "class EarlyStopping:\n",
+    "    def __init__(self, patience=10, verbose=False, delta=0.0001, save_path='./path/to/model/save'):\n",
+    "        self.patience = patience\n",
+    "        self.verbose = verbose\n",
+    "        self.delta = delta\n",
+    "        self.counter = 0\n",
+    "        self.best_score = None\n",
+    "        self.early_stop = False\n",
+    "        self.val_loss_min = float('inf')\n",
+    "        self.best_model = None\n",
+    "        self.save_path = save_path\n",
+    "        os.makedirs(save_path, exist_ok=True)\n",
+    "        \n",
+    "    def __call__(self, val_loss, model):\n",
+    "        score = -val_loss\n",
+    "\n",
+    "        if self.best_score is None:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "        elif score < self.best_score + self.delta:\n",
+    "            self.counter += 1\n",
+    "            if self.verbose:\n",
+    "                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
+    "            if self.counter >= self.patience:\n",
+    "                self.early_stop = True\n",
+    "        else:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "            self.counter = 0\n",
+    "\n",
+    "    def save_checkpoint(self, val_loss, model):\n",
+    "        if self.verbose:\n",
+    "            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')\n",
+    "        self.val_loss_min = val_loss\n",
+    "        self.best_model = model.state_dict()\n",
+    "        save_path = os.path.join(self.save_path, 'best_model.pth')\n",
+    "        torch.save(self.best_model, save_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c3fda74",
+   "metadata": {},
+   "source": [
+    "### Focal loss and reshape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fcf91db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ComplexFocalLoss(nn.Module):\n",
+    "    def __init__(self, alpha=1, gamma=2, reduction='mean'):\n",
+    "        super(ComplexFocalLoss, self).__init__()\n",
+    "        self.alpha = alpha\n",
+    "        self.gamma = gamma\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "        \n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction='none')\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction='none')\n",
+    "        \n",
+    "        real_pt = torch.exp(-real_BCE_loss)\n",
+    "        imag_pt = torch.exp(-imag_BCE_loss)\n",
+    "        \n",
+    "        real_F_loss = self.alpha * (1 - real_pt) ** self.gamma * real_BCE_loss\n",
+    "        imag_F_loss = self.alpha * (1 - imag_pt) ** self.gamma * imag_BCE_loss\n",
+    "\n",
+    "        if self.reduction == 'mean':\n",
+    "            return (torch.mean(real_F_loss) + torch.mean(imag_F_loss)) / 2\n",
+    "        elif self.reduction == 'sum':\n",
+    "            return torch.sum(real_F_loss) + torch.sum(imag_F_loss)\n",
+    "        else:\n",
+    "            return real_F_loss + imag_F_loss\n",
+    "\n",
+    "# Update the IoU calculation to handle complex values\n",
+    "def calculate_iou(pred, target, threshold=0.5):\n",
+    "    real_pred = (pred.real > threshold).float()\n",
+    "    imag_pred = (pred.imag > threshold).float()\n",
+    "    \n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred).float()\n",
+    "    \n",
+    "    intersection = (combined_pred * target).sum(dim=1)\n",
+    "    union = (combined_pred + target).sum(dim=1) - intersection\n",
+    "    iou = (intersection / union).mean().item()\n",
+    "    return iou\n",
+    "def reshape_to_2d(data):\n",
+    "    return data.view(-1, 1, 128, 128)  # Reshape to [batch, channels, height, width]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c97635b0",
+   "metadata": {},
+   "source": [
+    "### BCE Loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e8b2892",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# CV BCE Loss Function Definition\n",
+    "class ComplexValuedBCELoss(nn.Module):\n",
+    "    def __init__(self, reduction='mean'):\n",
+    "        super(ComplexValuedBCELoss, self).__init__()\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "\n",
+    "        # Calculate binary cross-entropy for both real and imaginary parts\n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction=self.reduction)\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction=self.reduction)\n",
+    "        \n",
+    "        # Combine the losses (you can adjust the weighting if necessary)\n",
+    "        combined_BCE_loss = (real_BCE_loss + imag_BCE_loss) / 2\n",
+    "        return combined_BCE_loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "64f4063c",
+   "metadata": {},
+   "source": [
+    "### Training from scratch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "66825110",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "device=\"cuda\"\n",
+    "def validate_model(model, valid_loader, criterion):\n",
+    "    model.eval()\n",
+    "    running_loss = 0.0\n",
+    "    iou_scores = []\n",
+    "    total_correct = 0\n",
+    "    total_samples = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, masks in tqdm(valid_loader, desc=\"Validating\"):\n",
+    "            inputs = reshape_to_2d(inputs).to(device)\n",
+    "            masks = masks.to(device)\n",
+    "            outputs = model(inputs)\n",
+    "            loss = criterion(outputs, masks)\n",
+    "            running_loss += loss.item()\n",
+    "\n",
+    "            # Calculate IoU\n",
+    "            iou = calculate_iou(outputs, masks, threshold=0.5)\n",
+    "            iou_scores.append(iou)\n",
+    "            \n",
+    "            # Calculate accuracy\n",
+    "            preds = (outputs.real > 0.5).float()\n",
+    "            correct = (preds == masks).float().sum()\n",
+    "            total_correct += correct.item()\n",
+    "            total_samples += masks.numel()\n",
+    "\n",
+    "    val_loss = running_loss / len(valid_loader)\n",
+    "    mean_iou = sum(iou_scores) / len(iou_scores)\n",
+    "    accuracy = total_correct / total_samples * 100\n",
+    "\n",
+    "    print(f'Validation Loss: {val_loss:.6f}')\n",
+    "    print(f'Validation Accuracy: {accuracy:.2f}%')\n",
+    "\n",
+    "    return val_loss, accuracy\n",
+    "\n",
+    "def train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.0001], num_epochs=50, patience=5):\n",
+    "    train_losses = []\n",
+    "    val_losses = []\n",
+    "    val_accuracies = []\n",
+    "    epoch_durations = []\n",
+    "    \n",
+    "    current_lr = initial_lr\n",
+    "    for lr in lr_steps:\n",
+    "        optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
+    "        early_stopping = EarlyStopping(patience=patience, verbose=True, delta=0.001)\n",
+    "        print(\"Current learning rate: \", lr)\n",
+    "        for epoch in range(num_epochs):\n",
+    "            epoch_start_time = time.time()\n",
+    "            \n",
+    "            model.train()\n",
+    "            running_loss = 0.0\n",
+    "            for inputs, masks in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} - Training\"):\n",
+    "                inputs = reshape_to_2d(inputs).to(device)\n",
+    "                masks = masks.to(device)\n",
+    "                outputs = model(inputs)\n",
+    "                loss = criterion(outputs, masks)\n",
+    "\n",
+    "                optimizer.zero_grad()\n",
+    "                loss.backward()\n",
+    "                optimizer.step()\n",
+    "\n",
+    "                running_loss += loss.item()\n",
+    "\n",
+    "            epoch_loss = running_loss / len(train_loader)\n",
+    "            train_losses.append(epoch_loss)\n",
+    "            print(f\"Training Loss: {epoch_loss:.6f}\")\n",
+    "            val_loss, val_accuracy = validate_model(model, valid_loader, criterion)\n",
+    "            val_losses.append(val_loss)\n",
+    "            val_accuracies.append(val_accuracy)\n",
+    "            early_stopping(val_loss, model)\n",
+    "\n",
+    "            if early_stopping.early_stop:\n",
+    "                print(\"Early stopping triggered\")\n",
+    "                break\n",
+    "\n",
+    "            epoch_duration = time.time() - epoch_start_time\n",
+    "            epoch_durations.append(epoch_duration)\n",
+    "        if early_stopping.best_model is not None:\n",
+    "            print(f\"Loading best model from lr {lr}\")\n",
+    "            model.load_state_dict(early_stopping.best_model)\n",
+    "        \n",
+    "    print(\"Training completed.\")\n",
+    "    print(\"Epoch durations:\", epoch_durations)\n",
+    "    return model, train_losses, val_losses, val_accuracies, epoch_durations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "621d28b3",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Initialize and train the ResNet-18 model\n",
+    "model = ComplexResNet18().to(device)\n",
+    "criterion = ComplexFocalLoss()\n",
+    "\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations =train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3)\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3838c1bc",
+   "metadata": {},
+   "source": [
+    "### Transfer Learning Load pretrained model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ac763e75",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Path to the pre-trained model weights\n",
+    "pretrained_model_path = \"path/to/model/save.pth\" #Change this model to trained model\n",
+    "device=\"cuda\"\n",
+    "# Initialize the model architecture\n",
+    "model = ComplexResNet18().to(device)\n",
+    "\n",
+    "# Load the pre-trained weights\n",
+    "checkpoint = torch.load(pretrained_model_path)\n",
+    "model.load_state_dict(checkpoint, strict=False)\n",
+    "\n",
+    "# Set all layers as trainable (if needed)\n",
+    "for param in model.parameters():\n",
+    "    param.requires_grad = True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1f877827",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Define a new criterion and optimizer for fine-tuning\n",
+    "# You may select between Focal Loss or BCE as your criterion\n",
+    "#criterion = ComplexValuedBCELoss()  # or ComplexValuedBCELoss()\n",
+    "criterion = ComplexFocalLoss()\n",
+    "# Use a smaller learning rate for fine-tuning\n",
+    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)\n",
+    "\n",
+    "# Train the model (fine-tuning)\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations= train_model(\n",
+    "    model, train_loader, valid_loader, criterion,\n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3\n",
+    ")\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f3784964",
+   "metadata": {},
+   "source": [
+    "### Plot Result and save the figures and json"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67a52e13",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Define save directory\n",
+    "save_dir = 'CMuSeNet_results/segmentation'\n",
+    "\n",
+    "# Create the directory if it doesn't exist\n",
+    "os.makedirs(save_dir, exist_ok=True)\n",
+    "\n",
+    "# Plot training loss\n",
+    "plt.figure()\n",
+    "plt.plot(range(1, len(train_losses) + 1), train_losses, label='Training Loss', color='blue')\n",
+    "plt.title('Training Loss')\n",
+    "plt.xlabel('Epoch')\n",
+    "plt.ylabel('Loss')\n",
+    "plt.legend()\n",
+    "\n",
+    "# Save the training loss figure as PNG and SVG\n",
+    "plt.savefig(os.path.join(save_dir, 'training_loss.png'))\n",
+    "plt.savefig(os.path.join(save_dir, 'training_loss.svg'))\n",
+    "\n",
+    "# Show the training loss plot\n",
+    "plt.show()\n",
+    "\n",
+    "# Plot validation accuracy\n",
+    "plt.figure()\n",
+    "plt.plot(range(1, len(val_accuracies) + 1), val_accuracies, label='Validation Accuracy', color='green')\n",
+    "plt.title('Validation Accuracy')\n",
+    "plt.xlabel('Epoch')\n",
+    "plt.ylabel('Accuracy')\n",
+    "plt.legend()\n",
+    "\n",
+    "# Save the validation accuracy figure as PNG and SVG\n",
+    "plt.savefig(os.path.join(save_dir, 'validation_accuracy.png'))\n",
+    "plt.savefig(os.path.join(save_dir, 'validation_accuracy.svg'))\n",
+    "\n",
+    "# Show the validation accuracy plot\n",
+    "plt.show()\n",
+    "\n",
+    "# Save the actual data to a JSON file\n",
+    "results = {\n",
+    "    \"train_losses\": train_losses,\n",
+    "    \"val_accuracies\": val_accuracies\n",
+    "}\n",
+    "\n",
+    "# Save JSON file\n",
+    "with open(os.path.join(save_dir, 'training_validation_results.json'), 'w') as f:\n",
+    "    json.dump(results, f)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "222069ae",
+   "metadata": {},
+   "source": [
+    "### BIG-RED Evaluation (Over entire dataset)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b178984",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "from tqdm import tqdm\n",
+    "# Create a DataLoader for the entire dataset\n",
+    "BATCH_SIZE = 64  # Adjust based on available memory\n",
+    "entire_dataset = WidebandSignalDataset(signal_ids=signal_dirs)  # Use all signals\n",
+    "entire_loader = DataLoader(entire_dataset, batch_size=BATCH_SIZE, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e6be59a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Path to the pre-trained model weights\n",
+    "pretrained_model_path = \"path/to/model/pretrained\" \n",
+    "device = \"cuda\" \n",
+    "\n",
+    "# Initialize the model architecture\n",
+    "model = ComplexResNet18().to(device)\n",
+    "\n",
+    "# Load the pre-trained weights\n",
+    "checkpoint = torch.load(pretrained_model_path, map_location=device)\n",
+    "model.load_state_dict(checkpoint, strict=False)\n",
+    "model.eval()\n",
+    "\n",
+    "# Function to evaluate accuracy\n",
+    "def evaluate_accuracy(model, data_loader):\n",
+    "    total_correct = 0\n",
+    "    total_samples = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, masks in tqdm(data_loader, desc=\"Evaluating on Entire Dataset\"):\n",
+    "            inputs = reshape_to_2d(inputs).to(device)\n",
+    "            masks = masks.to(device)\n",
+    "\n",
+    "            outputs = model(inputs)\n",
+    "            preds = (outputs.real > 0.5).float()\n",
+    "\n",
+    "            correct = (preds == masks).float().sum()\n",
+    "            total_correct += correct.item()\n",
+    "            total_samples += masks.numel()\n",
+    "\n",
+    "    accuracy = total_correct / total_samples * 100\n",
+    "    print(f\"Overall Accuracy on Entire Dataset: {accuracy:.2f}%\")\n",
+    "    return accuracy\n",
+    "\n",
+    "# Run the evaluation\n",
+    "overall_accuracy = evaluate_accuracy(model, entire_loader)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2a5a21b4",
+   "metadata": {},
+   "source": [
+    "### Function definitions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b223d9b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from tqdm import tqdm\n",
+    "import numpy as np\n",
+    "from collections import defaultdict\n",
+    "import torch.nn.functional as F\n",
+    "from scipy.optimize import linear_sum_assignment\n",
+    "from torch.utils.data import ConcatDataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f54736ea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the pre-trained model for evaluation\n",
+    "device = \"cuda\"\n",
+    "model_path = \"path/to/model/save.pth\"\n",
+    "model = resnet18_1D().to(device)\n",
+    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
+    "model.eval()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd5e7fee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "full_dataset = ConcatDataset([\n",
+    "    WidebandSignalDataset(signal_ids=train, return_snrs=True),\n",
+    "    WidebandSignalDataset(signal_ids=validation, return_snrs=True),\n",
+    "    WidebandSignalDataset(signal_ids=test, return_snrs=True)\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "173f9a8c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "full_loader = DataLoader(full_dataset, batch_size=64, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95f711d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def expand_true(array, distance=1):\n",
+    "    # Create kernel of appropriate size\n",
+    "    kernel = torch.ones((1, 1, distance * 2 + 1), device=array.device)\n",
+    "    array = array.unsqueeze(1).float()  # Add channel dimension\n",
+    "    result = F.conv1d(array, kernel, padding=distance)\n",
+    "    result = result.squeeze(1)  # Remove the extra dimension\n",
+    "    return result > 0\n",
+    "\n",
+    "def get_true_groups(tensor, device):\n",
+    "    assert tensor.dim() == 2, 'This function handles 2D tensor only'\n",
+    "    all_groups = []\n",
+    "    for i in range(tensor.size(0)):\n",
+    "        item = tensor[i]\n",
+    "        item = torch.cat([torch.tensor([False]).to(device), item, torch.tensor([False]).to(device)])\n",
+    "        diffs = item.float().diff()\n",
+    "        starts = (diffs == 1).nonzero(as_tuple=True)[0]\n",
+    "        ends = (diffs == -1).nonzero(as_tuple=True)[0] - 1\n",
+    "        groups = [(start.item(), end.item()) for start, end in zip(starts, ends)]\n",
+    "        all_groups.append(groups)\n",
+    "    return all_groups\n",
+    "\n",
+    "def calculate_iou(box1, box2):\n",
+    "    intersection = max(0, min(box1[1], box2[1]) - max(box1[0], box2[0]))\n",
+    "    union = max(box1[1], box2[1]) - min(box1[0], box2[0])\n",
+    "    return intersection / union if union != 0 else 0\n",
+    "\n",
+    "def match_targets(targets, preds):\n",
+    "    ious = []\n",
+    "    for target in targets:\n",
+    "        iou_targets = []\n",
+    "        for pred in preds:\n",
+    "            iou_targets.append(calculate_iou(target, pred))\n",
+    "        ious.append(iou_targets)\n",
+    "    cost_matrix = np.array(ious)\n",
+    "    row_ind, col_ind = linear_sum_assignment(-cost_matrix)\n",
+    "    return row_ind, col_ind\n",
+    "\n",
+    "def calculate_matched_ious(target_boxes, prediction_boxes, matching):\n",
+    "    ious = [0 for _ in target_boxes]\n",
+    "    matching_dict = dict(zip(*matching))\n",
+    "    for target_index, target_box in enumerate(target_boxes):\n",
+    "        if target_index in matching_dict:\n",
+    "            pred_index = matching_dict[target_index]\n",
+    "            if pred_index < len(prediction_boxes):\n",
+    "                box1 = target_box\n",
+    "                box2 = prediction_boxes[pred_index]\n",
+    "                ious[target_index] = calculate_iou(box1, box2)\n",
+    "    return ious\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "40ec3d9f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def evaluate(predictor, data_loader, device=\"cuda\"):\n",
+    "    iou_thresholds = [0.5, 0.7, 0.9]\n",
+    "    snr_metrics = defaultdict(lambda: {\n",
+    "        \"iou_sum\": 0.0,\n",
+    "        \"iou_count\": 0,\n",
+    "        \"recall_counts\": defaultdict(int),\n",
+    "        \"total_samples\": defaultdict(int),\n",
+    "        \"correct_pixels\": 0,\n",
+    "        \"total_pixels\": 0\n",
+    "    })\n",
+    "    total_iou_sum, total_iou_count = 0.0, 0\n",
+    "    total_correct_pixels, total_total_pixels = 0, 0\n",
+    "    total_recall_counts = defaultdict(int)\n",
+    "    total_samples = defaultdict(int)\n",
+    "\n",
+    "    for batch in tqdm(data_loader, desc=\"Evaluating\"):\n",
+    "        if len(batch) == 3:\n",
+    "            inputs, masks, snrs_in_batch = batch\n",
+    "        else:\n",
+    "            inputs, masks = batch\n",
+    "            snrs_in_batch = [0] * len(inputs)  # Default SNR if not provided\n",
+    "\n",
+    "        inputs = inputs.to(device)\n",
+    "        masks = masks.to(device)\n",
+    "        outputs = predictor(inputs)\n",
+    "\n",
+    "        for i in range(len(inputs)):\n",
+    "            mask = masks[i]\n",
+    "            output = outputs[i]\n",
+    "\n",
+    "            # Resize output to match mask shape if necessary\n",
+    "            if output.numel() != mask.numel():\n",
+    "                output = output.expand_as(mask) if output.numel() == 1 else output.reshape_as(mask)\n",
+    "\n",
+    "            thresholded_output = (output >= 0.5).float()\n",
+    "\n",
+    "            correct_pixels = (thresholded_output == mask).sum().item()\n",
+    "            total_pixels = mask.numel()\n",
+    "            total_correct_pixels += correct_pixels\n",
+    "            total_total_pixels += total_pixels\n",
+    "\n",
+    "            # Get SNR value and round it to the nearest integer\n",
+    "            snr = snrs_in_batch[i]\n",
+    "            if isinstance(snr, torch.Tensor):\n",
+    "                snr = snr.item()\n",
+    "            snr = int(round(snr))  # Round SNR to the nearest integer\n",
+    "\n",
+    "            snr_metrics[snr][\"correct_pixels\"] += correct_pixels\n",
+    "            snr_metrics[snr][\"total_pixels\"] += total_pixels\n",
+    "\n",
+    "            target_boxes = get_true_groups(mask.unsqueeze(0), device=device)[0]\n",
+    "            pred_boxes = get_true_groups(thresholded_output.unsqueeze(0), device=device)[0]\n",
+    "            if not target_boxes or not pred_boxes:\n",
+    "                continue\n",
+    "            matching = match_targets(target_boxes, pred_boxes)\n",
+    "            matched_ious = calculate_matched_ious(target_boxes, pred_boxes, matching)\n",
+    "\n",
+    "            snr_metrics[snr][\"iou_sum\"] += sum(matched_ious)\n",
+    "            snr_metrics[snr][\"iou_count\"] += len(matched_ious)\n",
+    "            total_iou_sum += sum(matched_ious)\n",
+    "            total_iou_count += len(matched_ious)\n",
+    "\n",
+    "            for th in iou_thresholds:\n",
+    "                true_positives = sum(1 for iou in matched_ious if iou >= th)\n",
+    "                snr_metrics[snr][\"recall_counts\"][th] += true_positives\n",
+    "                snr_metrics[snr][\"total_samples\"][th] += len(target_boxes)\n",
+    "                total_recall_counts[th] += true_positives\n",
+    "                total_samples[th] += len(target_boxes)\n",
+    "\n",
+    "    # Calculate overall metrics\n",
+    "    overall_accuracy = (total_correct_pixels / total_total_pixels) * 100 if total_total_pixels > 0 else 0\n",
+    "    overall_iou = total_iou_sum / total_iou_count if total_iou_count > 0 else 0\n",
+    "    overall_recall = {\n",
+    "        th: total_recall_counts[th] / total_samples[th] if total_samples[th] > 0 else 0\n",
+    "        for th in iou_thresholds\n",
+    "    }\n",
+    "\n",
+    "    # Print overall results\n",
+    "    print(f\"Overall Accuracy: {overall_accuracy:.2f}%\")\n",
+    "    print(f\"Overall IoU Score: {overall_iou:.4f}\")\n",
+    "    for th in iou_thresholds:\n",
+    "        print(f\"Recall at threshold {th}: {overall_recall[th]:.4f}\")\n",
+    "\n",
+    "    # Print per-SNR results\n",
+    "    for snr in sorted(snr_metrics.keys()):\n",
+    "        metrics = snr_metrics[snr]\n",
+    "        snr_accuracy = (metrics[\"correct_pixels\"] / metrics[\"total_pixels\"]) * 100 if metrics[\"total_pixels\"] > 0 else 0\n",
+    "        snr_iou = metrics[\"iou_sum\"] / metrics[\"iou_count\"] if metrics[\"iou_count\"] > 0 else 0\n",
+    "        print(f\"SNR: {snr} dB - Accuracy: {snr_accuracy:.2f}%\")\n",
+    "        print(f\"   IoU: {snr_iou:.4f}\")\n",
+    "        for th in iou_thresholds:\n",
+    "            recall = metrics[\"recall_counts\"][th] / metrics[\"total_samples\"][th] if metrics[\"total_samples\"][th] > 0 else 0\n",
+    "            print(f\"   Recall at threshold {th}: {recall:.4f}\")\n",
+    "\n",
+    "    return snr_metrics\n",
+    "\n",
+    "\n",
+    "def model_predictor(signals):\n",
+    "    # Use the already loaded model and apply thresholding\n",
+    "    return expand_true(model(signals) > 0.5)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7d3aed7",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Run evaluation on the full dataset\n",
+    "snr_metrics = evaluate(model_predictor, full_loader, device=device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2fd3ba0e",
+   "metadata": {},
+   "source": [
+    "### Save and plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aef69113",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "def save_results_and_plot(snr_metrics, save_path):\n",
+    "    \"\"\"\n",
+    "    Saves evaluation results to a JSON file and generates plots for Accuracy, IoU, and Recall vs. SNR.\n",
+    "    Sets x-axis limits to range from -9 dB to 12 dB to eliminate blank space on the right.\n",
+    "\n",
+    "    Args:\n",
+    "        snr_metrics (dict): The evaluation results obtained from the evaluate function.\n",
+    "        save_path (str): The directory path where results and plots will be saved.\n",
+    "\n",
+    "    Outputs:\n",
+    "        - evaluation_results.json\n",
+    "        - accuracy_vs_snr.png and .svg\n",
+    "        - iou_vs_snr.png and .svg\n",
+    "        - recall_vs_snr.png and .svg\n",
+    "    \"\"\"\n",
+    "    # Ensure the directory exists\n",
+    "    os.makedirs(save_path, exist_ok=True)\n",
+    "    \n",
+    "    # Extract data from snr_metrics\n",
+    "    snr_list = sorted(snr_metrics.keys())\n",
+    "    accuracy_list = []\n",
+    "    iou_list = []\n",
+    "    recall_05 = []\n",
+    "    recall_07 = []\n",
+    "    recall_09 = []\n",
+    "    \n",
+    "    # Prepare data for JSON serialization\n",
+    "    json_data = {}\n",
+    "    \n",
+    "    for snr in snr_list:\n",
+    "        metrics = snr_metrics[snr]\n",
+    "        snr_accuracy = (metrics[\"correct_pixels\"] / metrics[\"total_pixels\"]) * 100 if metrics[\"total_pixels\"] > 0 else 0\n",
+    "        snr_iou = metrics[\"iou_sum\"] / metrics[\"iou_count\"] if metrics[\"iou_count\"] > 0 else 0\n",
+    "        recall_at_05 = metrics[\"recall_counts\"][0.5] / metrics[\"total_samples\"][0.5] if metrics[\"total_samples\"][0.5] > 0 else 0\n",
+    "        recall_at_07 = metrics[\"recall_counts\"][0.7] / metrics[\"total_samples\"][0.7] if metrics[\"total_samples\"][0.7] > 0 else 0\n",
+    "        recall_at_09 = metrics[\"recall_counts\"][0.9] / metrics[\"total_samples\"][0.9] if metrics[\"total_samples\"][0.9] > 0 else 0\n",
+    "\n",
+    "        # Append to lists for plotting\n",
+    "        accuracy_list.append(snr_accuracy)\n",
+    "        iou_list.append(snr_iou)\n",
+    "        recall_05.append(recall_at_05)\n",
+    "        recall_07.append(recall_at_07)\n",
+    "        recall_09.append(recall_at_09)\n",
+    "\n",
+    "        # Prepare data for JSON\n",
+    "        json_data[snr] = {\n",
+    "            \"accuracy\": snr_accuracy,\n",
+    "            \"iou\": snr_iou,\n",
+    "            \"recall\": {\n",
+    "                \"0.5\": recall_at_05,\n",
+    "                \"0.7\": recall_at_07,\n",
+    "                \"0.9\": recall_at_09,\n",
+    "            }\n",
+    "        }\n",
+    "    \n",
+    "    # Save json_data to JSON file\n",
+    "    json_file_path = os.path.join(save_path, 'evaluation_results.json')\n",
+    "    with open(json_file_path, 'w') as json_file:\n",
+    "        json.dump(json_data, json_file, indent=4)\n",
+    "    \n",
+    "    # Plot Accuracy vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, accuracy_list, marker='o', label='Accuracy')\n",
+    "    plt.title('Accuracy vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('Accuracy (%)')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    plt.xlim(-9, 12)\n",
+    "    \n",
+    "    # Save the plot\n",
+    "    accuracy_png_path = os.path.join(save_path, 'accuracy_vs_snr.png')\n",
+    "    accuracy_svg_path = os.path.join(save_path, 'accuracy_vs_snr.svg')\n",
+    "    plt.savefig(accuracy_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(accuracy_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n",
+    "    \n",
+    "    # Plot IoU vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, iou_list, marker='o', color='orange', label='IoU')\n",
+    "    plt.title('IoU vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('IoU')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    plt.xlim(-9, 12)\n",
+    "    \n",
+    "    # Save the plot\n",
+    "    iou_png_path = os.path.join(save_path, 'iou_vs_snr.png')\n",
+    "    iou_svg_path = os.path.join(save_path, 'iou_vs_snr.svg')\n",
+    "    plt.savefig(iou_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(iou_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n",
+    "    \n",
+    "    # Plot Recall at Different IoU Thresholds vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, recall_05, marker='o', label='Recall @ IoU 0.5')\n",
+    "    plt.plot(snr_list, recall_07, marker='s', label='Recall @ IoU 0.7')\n",
+    "    plt.plot(snr_list, recall_09, marker='^', label='Recall @ IoU 0.9')\n",
+    "    plt.title('Recall at Different IoU Thresholds vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('Recall')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    plt.xlim(-9, 12)\n",
+    "    \n",
+    "    # Save the plot\n",
+    "    recall_png_path = os.path.join(save_path, 'recall_vs_snr.png')\n",
+    "    recall_svg_path = os.path.join(save_path, 'recall_vs_snr.svg')\n",
+    "    plt.savefig(recall_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(recall_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c9595d5e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Assuming snr_metrics is the output from the evaluate function\n",
+    "# Set the save path\n",
+    "save_path = 'CMuSeNet_BIGRED_results'\n",
+    "\n",
+    "# Call the function\n",
+    "save_results_and_plot(snr_metrics, save_path)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

CMuSeNet_Indoor_OTA.ipynb

@@ -0,0 +1,1658 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b5007b71",
+   "metadata": {},
+   "source": [
+    "### Initialization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e6b1226",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Initialization block\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import json\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "import math\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
+    "\n",
+    "STFT_LENGTH = 16 * 1024\n",
+    "DATA_DIR = Path(\"/data/OTA_reduced/\")\n",
+    "SAMPLE_RATE = 20e6\n",
+    "MODULATIONS = [\"QPSK\", \"BPSK\", \"2-FSK\"]\n",
+    "MODULATION_LABELS = {j: i for i, j in enumerate(MODULATIONS)}\n",
+    "NUMBER_OF_MODULATIONS = len(MODULATIONS)\n",
+    "MASK_SIZE = int(STFT_LENGTH)\n",
+    "\n",
+    "from matplotlib.mlab import psd as apply_psd\n",
+    "\n",
+    "def calc_sig_power(signal, meta, noise_power=-132.065):\n",
+    "    \n",
+    "    noise_floor_linear = 10 ** (noise_power / 10)\n",
+    "    (psd, frequencies) = apply_psd(signal, Fs=SAMPLE_RATE, NFFT=1024)\n",
+    "\n",
+    "\n",
+    "    signal_position = []\n",
+    "\n",
+    "    body = meta[\"body\"]\n",
+    "    device = meta[\"client_id\"]\n",
+    "    bandwidth, frequency_offset = body[\"bandwidth\"] + 20e3, body[\"frequency_offset\"]\n",
+    "\n",
+    "    \n",
+    "    below_freq = frequency_offset-bandwidth/2\n",
+    "    upper_freq = frequency_offset+bandwidth/2\n",
+    "    sum_power_dbs = 0\n",
+    "    freq_count = 0\n",
+    "    \n",
+    "    for idx, (power, freq) in enumerate(zip(psd, frequencies)):\n",
+    "        if below_freq <= freq <= upper_freq:\n",
+    "            freq_count+=1\n",
+    "            sum_power_dbs+=(power)\n",
+    "    return sum_power_dbs\n",
+    "\n",
+    "# noise_power is measured from noise signal collection\n",
+    "def calc_snr(signal_power, noise_power=-132.065):\n",
+    "    noise_floor_linear = 10 ** (noise_power / 10)\n",
+    "    snr_linear = signal_power / (noise_floor_linear * 1024)\n",
+    "    \n",
+    "    snr_db = 10 * np.log10(snr_linear)\n",
+    "    \n",
+    "    return round(snr_db)\n",
+    "\n",
+    "def convert_metadata_format_real_to_simulated(signal, metadata):\n",
+    "    name_mapping = {\"2FSK\": \"2-FSK\"}\n",
+    "    return [\n",
+    "        {\n",
+    "            \"fc\": body[\"frequency_offset\"], \n",
+    "            \"bw\": body[\"bandwidth\"] + 20e3,\n",
+    "            \"mod\": name_mapping.get(body[\"modulation\"], body[\"modulation\"]),\n",
+    "            \"snr\": calc_snr(calc_sig_power(signal, meta))\n",
+    "        } for meta in metadata if (body := meta[\"body\"])\n",
+    "    ]\n",
+    "\n",
+    "def load_data(signal_id, load_metadata_only=False):\n",
+    "    if not load_metadata_only:\n",
+    "        signal_path = DATA_DIR / str(signal_id) / \"data.npy\"\n",
+    "        if not signal_path.exists():\n",
+    "            raise FileNotFoundError(f\"Signal file {signal_path} not found.\")\n",
+    "        signal = np.load(signal_path)\n",
+    "    else:\n",
+    "        signal = None\n",
+    "    with open(DATA_DIR / str(signal_id) / \"meta-data.json\") as f:\n",
+    "        meta = json.load(f)\n",
+    "        if isinstance(meta, dict):\n",
+    "            meta = [meta]\n",
+    "    return signal, convert_metadata_format_real_to_simulated(signal, meta)\n",
+    "\n",
+    "\n",
+    "    \n",
+    "def _get_all_numbered_dirs(root_dir):\n",
+    "    dirs = []\n",
+    "    for directory in root_dir.iterdir():\n",
+    "        dirs.append(int(directory.name))\n",
+    "    dirs.sort()\n",
+    "    return dirs\n",
+    "        \n",
+    "        \n",
+    "def process_metadata(metadata):\n",
+    "    scaled_metadata =  [\n",
+    "        {\n",
+    "            \"position\": (SAMPLE_RATE/2 + i['fc'], i['bw']),\n",
+    "            \"mod\": i[\"mod\"],\n",
+    "            \"snr\": i[\"snr\"],\n",
+    "            \"bw\": int(i['bw'])\n",
+    "        }\n",
+    "        for i in metadata\n",
+    "    ]\n",
+    "    return scaled_metadata\n",
+    "\n",
+    "\n",
+    "def process_signal(signal):\n",
+    "    signal = signal[:STFT_LENGTH]\n",
+    "\n",
+    "    signal = np.fft.fft(signal)\n",
+    "    signal = np.fft.fftshift(signal)\n",
+    "    signal /= np.max(np.abs(signal))\n",
+    "    return signal"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "440b802c",
+   "metadata": {},
+   "source": [
+    "### Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31bc3770",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class WidebandSignalDataset(torch.utils.data.Dataset):\n",
+    "    def __init__(self, signal_ids, mask_size=MASK_SIZE, return_snrs=False):\n",
+    "        self.mask_size = mask_size\n",
+    "        self.signal_ids = signal_ids\n",
+    "        self.return_snrs = return_snrs\n",
+    "        self.snrs = []\n",
+    "        loaded_data = []\n",
+    "        \n",
+    "        for signal_id in tqdm(self.signal_ids):\n",
+    "            loaded_data.append(self.process_signal(signal_id))\n",
+    "            \n",
+    "        self.loaded_data = loaded_data\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.signal_ids)\n",
+    "\n",
+    "    def __getitem__(self, index):\n",
+    "        if self.return_snrs:\n",
+    "            signal, masks, snr = self.loaded_data[index]\n",
+    "        else:\n",
+    "            signal, masks = self.loaded_data[index]\n",
+    "\n",
+    "        # Ensure `signal` is complex and `masks` is real-valued\n",
+    "        if not isinstance(signal, torch.Tensor):\n",
+    "            signal = torch.from_numpy(signal).type(torch.complex64)\n",
+    "        if not isinstance(masks, torch.Tensor):\n",
+    "            masks = torch.from_numpy(masks).type(torch.FloatTensor)\n",
+    "\n",
+    "        if self.return_snrs:\n",
+    "            if not isinstance(snr, torch.Tensor):\n",
+    "                snr = torch.tensor(snr).type(torch.FloatTensor)\n",
+    "            return signal, masks, snr\n",
+    "        else:\n",
+    "            return signal, masks\n",
+    "\n",
+    "    def process_signal(self, signal_id):\n",
+    "        # Load data and metadata\n",
+    "        signal, metadata = load_data(signal_id)\n",
+    "        \n",
+    "        # Process the metadata and create masks\n",
+    "        scaled_metadata = process_metadata(metadata)\n",
+    "        snrs = [meta['snr'] for meta in scaled_metadata]\n",
+    "        average_snr = sum(snrs) / len(snrs) if snrs else 0\n",
+    "        \n",
+    "        # Convert signal to complex format and normalize it\n",
+    "        signal = process_signal(signal)  # `process_signal` should return np.ndarray (complex)\n",
+    "        signal = torch.from_numpy(signal).type(torch.complex64)  # Convert to complex tensor\n",
+    "        \n",
+    "        # Generate binary mask for each frequency segment\n",
+    "        masks = np.zeros(self.mask_size, dtype=np.float32)\n",
+    "        scale_ratio = self.mask_size / SAMPLE_RATE\n",
+    "        for meta in scaled_metadata:\n",
+    "            f, b = meta['position']\n",
+    "            x1 = math.floor((f - b / 2) * scale_ratio)\n",
+    "            x2 = math.ceil((f + b / 2) * scale_ratio)\n",
+    "            masks[x1:x2] = 1\n",
+    "        \n",
+    "        if self.return_snrs:\n",
+    "            return signal, masks, average_snr\n",
+    "        else:\n",
+    "            return signal, masks\n",
+    "\n",
+    "\n",
+    "# Train test split 80 - 10 - 10\n",
+    "train, test, validation = [], [], [] \n",
+    "total_signals = len([i for i in DATA_DIR.iterdir()])\n",
+    "for index, signal in enumerate(_get_all_numbered_dirs(DATA_DIR)):\n",
+    "    if index <= 0.80 * total_signals:\n",
+    "        train.append(signal)\n",
+    "    elif index <= 0.9 * total_signals:\n",
+    "        validation.append(signal)\n",
+    "    else:\n",
+    "        test.append(signal)\n",
+    "            \n",
+    "print(\"Train\", len(train))\n",
+    "print(\"Validation\", len(validation))\n",
+    "print(\"Test\", len(test))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3e74df1a",
+   "metadata": {},
+   "source": [
+    "### Check if complex value"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23f75344",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test_single_signal_loading(signal_id):\n",
+    "    # Load a single signal and process it\n",
+    "    signal, metadata = load_data(signal_id)\n",
+    "    \n",
+    "    # Process the signal: Apply any necessary preprocessing, and convert to complex format\n",
+    "    processed_signal = process_signal(signal)  # This should return a complex np.ndarray\n",
+    "    complex_signal = torch.from_numpy(processed_signal).type(torch.complex64)\n",
+    "    \n",
+    "    # Check if the signal is complex\n",
+    "    print(\"Loaded Signal ID:\", signal_id)\n",
+    "    print(\"Signal Type:\", complex_signal.dtype)\n",
+    "    print(\"Signal Shape:\", complex_signal.shape)\n",
+    "    \n",
+    "    # Generate the mask as you would in WidebandSignalDataset\n",
+    "    scaled_metadata = process_metadata(metadata)\n",
+    "    masks = np.zeros(MASK_SIZE, dtype=np.float32)\n",
+    "    scale_ratio = MASK_SIZE / SAMPLE_RATE\n",
+    "    for meta in scaled_metadata:\n",
+    "        f, b = meta['position']\n",
+    "        x1 = math.floor((f - b / 2) * scale_ratio)\n",
+    "        x2 = math.ceil((f + b / 2) * scale_ratio)\n",
+    "        masks[x1:x2] = 1\n",
+    "\n",
+    "    # Convert mask to tensor\n",
+    "    mask_tensor = torch.from_numpy(masks).type(torch.FloatTensor)\n",
+    "\n",
+    "    # Output information about the mask\n",
+    "    print(\"Mask Shape:\", mask_tensor.shape)\n",
+    "    print(\"Mask Type:\", mask_tensor.dtype)\n",
+    "    \n",
+    "    return complex_signal, mask_tensor\n",
+    "\n",
+    "# Test with a specific signal_id (replace with an actual ID from your data)\n",
+    "test_signal_id = train[0]  # Assuming `train` list contains valid signal IDs\n",
+    "complex_signal, mask_tensor = test_single_signal_loading(test_signal_id)\n",
+    "\n",
+    "# Optional: Check a sample value to confirm it's complex\n",
+    "print(\"Sample value from signal tensor:\", complex_signal[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1cec9c6e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataset = WidebandSignalDataset(signal_ids=train)\n",
+    "validation_dataset = WidebandSignalDataset(signal_ids=validation)\n",
+    "test_dataset = WidebandSignalDataset(signal_ids=test)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e0900d4e",
+   "metadata": {},
+   "source": [
+    "### Check SNR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2fbee106",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# For Train Dataset\n",
+    "train_snrs = train_dataset.snrs\n",
+    "\n",
+    "# Plot Histogram of SNRs in Train Dataset\n",
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.hist(train_snrs, bins=range(int(min(train_snrs)), int(max(train_snrs)) + 1), edgecolor='black')\n",
+    "plt.title('Histogram of SNRs in Train Dataset')\n",
+    "plt.xlabel('SNR (dB)')\n",
+    "plt.ylabel('Number of Samples')\n",
+    "plt.grid(True)\n",
+    "plt.show()\n",
+    "\n",
+    "# Print SNR Range\n",
+    "print('Train Dataset SNR range: {} dB to {} dB'.format(min(train_snrs), max(train_snrs)))\n",
+    "\n",
+    "# For Validation Dataset\n",
+    "validation_snrs = validation_dataset.snrs\n",
+    "\n",
+    "# Plot Histogram of SNRs in Validation Dataset\n",
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.hist(validation_snrs, bins=range(int(min(validation_snrs)), int(max(validation_snrs)) + 1), edgecolor='black')\n",
+    "plt.title('Histogram of SNRs in Validation Dataset')\n",
+    "plt.xlabel('SNR (dB)')\n",
+    "plt.ylabel('Number of Samples')\n",
+    "plt.grid(True)\n",
+    "plt.show()\n",
+    "\n",
+    "# Print SNR Range\n",
+    "print('Validation Dataset SNR range: {} dB to {} dB'.format(min(validation_snrs), max(validation_snrs)))\n",
+    "\n",
+    "# For Test Dataset\n",
+    "test_snrs = test_dataset.snrs\n",
+    "\n",
+    "# Plot Histogram of SNRs in Validation Dataset\n",
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.hist(test_snrs, bins=range(int(min(test_snrs)), int(max(test_snrs)) + 1), edgecolor='black')\n",
+    "plt.title('Histogram of SNRs in Test Dataset')\n",
+    "plt.xlabel('SNR (dB)')\n",
+    "plt.ylabel('Number of Samples')\n",
+    "plt.grid(True)\n",
+    "plt.show()\n",
+    "\n",
+    "# Print SNR Range\n",
+    "print('Validation Dataset SNR range: {} dB to {} dB'.format(min(test_snrs), max(test_snrs)))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "637ae774",
+   "metadata": {},
+   "source": [
+    "### Batch Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9af2450",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 64  # Updated batch size\n",
+    "\n",
+    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
+    "valid_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False)\n",
+    "\n",
+    "print(\"Train labels shape:\", len(train_dataset))\n",
+    "print(\"Validation labels shape:\", len(validation_dataset))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a8e09e4",
+   "metadata": {},
+   "source": [
+    "### Early Stop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24f79a24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "class EarlyStopping:\n",
+    "    def __init__(self, patience=10, verbose=False, delta=0.0001, save_path='./path/to/model/save'):\n",
+    "        self.patience = patience\n",
+    "        self.verbose = verbose\n",
+    "        self.delta = delta\n",
+    "        self.counter = 0\n",
+    "        self.best_score = None\n",
+    "        self.early_stop = False\n",
+    "        self.val_loss_min = float('inf')\n",
+    "        self.best_model = None\n",
+    "        self.save_path = save_path\n",
+    "        os.makedirs(save_path, exist_ok=True)\n",
+    "        \n",
+    "    def __call__(self, val_loss, model):\n",
+    "        score = -val_loss\n",
+    "\n",
+    "        if self.best_score is None:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "        elif score < self.best_score + self.delta:\n",
+    "            self.counter += 1\n",
+    "            if self.verbose:\n",
+    "                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
+    "            if self.counter >= self.patience:\n",
+    "                self.early_stop = True\n",
+    "        else:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "            self.counter = 0\n",
+    "\n",
+    "    def save_checkpoint(self, val_loss, model):\n",
+    "        if self.verbose:\n",
+    "            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')\n",
+    "        self.val_loss_min = val_loss\n",
+    "        self.best_model = model.state_dict()\n",
+    "        save_path = os.path.join(self.save_path, 'best_model.pth')\n",
+    "        torch.save(self.best_model, save_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c3fda74",
+   "metadata": {},
+   "source": [
+    "### Reshape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fcf91db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn as nn\n",
+    "import complexPyTorch.complexLayers as cplx\n",
+    "import torch.nn.functional as F\n",
+    "import torch\n",
+    "\n",
+    "def reshape_to_2d(data):\n",
+    "    return data.view(-1, 1, 128, 128)  # Reshape to [batch, channels, height, width]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7d7562c",
+   "metadata": {},
+   "source": [
+    "### Complex IoU"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76b9d084",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def calculate_iou(pred, target, threshold=0.5):\n",
+    "    real_pred = (pred.real > threshold).float()\n",
+    "    imag_pred = (pred.imag > threshold).float()\n",
+    "    \n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred).float()\n",
+    "    \n",
+    "    intersection = (combined_pred * target).sum(dim=1)\n",
+    "    union = (combined_pred + target).sum(dim=1) - intersection\n",
+    "    iou = (intersection / union).mean().item()\n",
+    "    return iou"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "64f4063c",
+   "metadata": {},
+   "source": [
+    "### Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "66825110",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "\n",
+    "def validate_model(model, valid_loader, criterion):\n",
+    "    model.eval()\n",
+    "    running_loss = 0.0\n",
+    "    iou_scores = []\n",
+    "    total_correct = 0\n",
+    "    total_samples = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, masks in tqdm(valid_loader, desc=\"Validating\"):\n",
+    "            inputs = reshape_to_2d(inputs).to(device)\n",
+    "            masks = masks.to(device)\n",
+    "            outputs = model(inputs)\n",
+    "            loss = criterion(outputs, masks)\n",
+    "            running_loss += loss.item()\n",
+    "\n",
+    "            # Calculate IoU\n",
+    "            iou = calculate_iou(outputs, masks, threshold=0.5)\n",
+    "            iou_scores.append(iou)\n",
+    "            \n",
+    "            # Calculate accuracy\n",
+    "            preds = (outputs.real > 0.5).float()\n",
+    "            correct = (preds == masks).float().sum()\n",
+    "            total_correct += correct.item()\n",
+    "            total_samples += masks.numel()\n",
+    "\n",
+    "    val_loss = running_loss / len(valid_loader)\n",
+    "    mean_iou = sum(iou_scores) / len(iou_scores)\n",
+    "    accuracy = total_correct / total_samples * 100\n",
+    "\n",
+    "    print(f'Validation Loss: {val_loss:.6f}')\n",
+    "    print(f'Validation Accuracy: {accuracy:.2f}%')\n",
+    "\n",
+    "    return val_loss, accuracy\n",
+    "\n",
+    "def train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.000001], num_epochs=50, patience=5):\n",
+    "    train_losses = []\n",
+    "    val_losses = []\n",
+    "    val_accuracies = []\n",
+    "    epoch_durations = []\n",
+    "    \n",
+    "    current_lr = initial_lr\n",
+    "    for lr in lr_steps:\n",
+    "        optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
+    "        early_stopping = EarlyStopping(patience=patience, verbose=True, delta=0.001)\n",
+    "        print(\"Current learning rate: \", lr)\n",
+    "        for epoch in range(num_epochs):\n",
+    "            epoch_start_time = time.time()\n",
+    "            \n",
+    "            model.train()\n",
+    "            running_loss = 0.0\n",
+    "            for inputs, masks in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} - Training\"):\n",
+    "                inputs = reshape_to_2d(inputs).to(device)\n",
+    "                masks = masks.to(device)\n",
+    "                outputs = model(inputs)\n",
+    "                loss = criterion(outputs, masks)\n",
+    "\n",
+    "                optimizer.zero_grad()\n",
+    "                loss.backward()\n",
+    "                optimizer.step()\n",
+    "\n",
+    "                running_loss += loss.item()\n",
+    "\n",
+    "            epoch_loss = running_loss / len(train_loader)\n",
+    "            train_losses.append(epoch_loss)\n",
+    "            print(f\"Training Loss: {epoch_loss:.6f}\")\n",
+    "            \n",
+    "            val_loss, val_accuracy = validate_model(model, valid_loader, criterion)\n",
+    "            val_losses.append(val_loss)\n",
+    "            val_accuracies.append(val_accuracy)\n",
+    "            early_stopping(val_loss, model)\n",
+    "\n",
+    "            if early_stopping.early_stop:\n",
+    "                print(\"Early stopping triggered\")\n",
+    "                break\n",
+    "\n",
+    "            epoch_duration = time.time() - epoch_start_time\n",
+    "            epoch_durations.append(epoch_duration)\n",
+    "        if early_stopping.best_model is not None:\n",
+    "            print(f\"Loading best model from lr {lr}\")\n",
+    "            model.load_state_dict(early_stopping.best_model)\n",
+    "        \n",
+    "    print(\"Training completed.\")\n",
+    "    print(\"Epoch durations:\", epoch_durations)\n",
+    "    return model, train_losses, val_losses, val_accuracies, epoch_durations"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b80cb51",
+   "metadata": {},
+   "source": [
+    "### ResNet-18"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d208cb9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import complexPyTorch.complexLayers as cplx\n",
+    "from typing import Optional, Callable, Type, Union, List\n",
+    "import torch.nn.functional as F\n",
+    "from torch import Tensor\n",
+    "\n",
+    "def conv3x3(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"3x3 convolution with padding\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
+    "\n",
+    "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"1x1 convolution\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
+    "\n",
+    "class BasicBlock(nn.Module):\n",
+    "    expansion = 1\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(BasicBlock, self).__init__()\n",
+    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.conv2 = conv3x3(planes, planes)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class Bottleneck(nn.Module):\n",
+    "    expansion = 4\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(Bottleneck, self).__init__()\n",
+    "        self.conv1 = conv1x1(inplanes, planes)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv2 = conv3x3(planes, planes, stride)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv3 = conv1x1(planes, planes * self.expansion)\n",
+    "        self.bn3 = cplx.ComplexBatchNorm2d(planes * self.expansion)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv3(out)\n",
+    "        out = self.bn3(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class ComplexResNet(nn.Module):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        block: Type[Union[BasicBlock, Bottleneck]],\n",
+    "        layers: List[int],\n",
+    "        num_classes: int = STFT_LENGTH,\n",
+    "        zero_init_residual: bool = False,\n",
+    "        groups: int = 1,\n",
+    "        width_per_group: int = 64,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(ComplexResNet, self).__init__()\n",
+    "        if norm_layer is None:\n",
+    "            norm_layer = cplx.ComplexBatchNorm2d\n",
+    "        self._norm_layer = norm_layer\n",
+    "\n",
+    "        self.inplanes = 64\n",
+    "        self.dilation = 1\n",
+    "\n",
+    "        self.groups = groups\n",
+    "        self.base_width = width_per_group\n",
+    "        self.conv1 = cplx.ComplexConv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)\n",
+    "        self.bn1 = norm_layer(self.inplanes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
+    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
+    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
+    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
+    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
+    "        self.fc = cplx.ComplexLinear(512 * block.expansion, num_classes)\n",
+    "        self.sigmoid = cplx.ComplexSigmoid()\n",
+    "\n",
+    "    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int, stride: int = 1) -> nn.Sequential:\n",
+    "        norm_layer = self._norm_layer\n",
+    "        downsample = None\n",
+    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
+    "            downsample = nn.Sequential(\n",
+    "                conv1x1(self.inplanes, planes * block.expansion, stride),\n",
+    "                norm_layer(planes * block.expansion),\n",
+    "            )\n",
+    "\n",
+    "        layers = []\n",
+    "        layers.append(block(self.inplanes, planes, stride, downsample, norm_layer))\n",
+    "        self.inplanes = planes * block.expansion\n",
+    "        for _ in range(1, blocks):\n",
+    "            layers.append(block(self.inplanes, planes, norm_layer=norm_layer))\n",
+    "\n",
+    "        return nn.Sequential(*layers)\n",
+    "\n",
+    "    def _forward_impl(self, x: Tensor) -> Tensor:\n",
+    "        x = self.conv1(x)\n",
+    "        x = self.bn1(x)\n",
+    "        x = self.relu(x)\n",
+    "\n",
+    "        x = self.layer1(x)\n",
+    "        x = self.layer2(x)\n",
+    "        x = self.layer3(x)\n",
+    "        x = self.layer4(x)\n",
+    "\n",
+    "        x = self.avgpool(x)\n",
+    "        x = torch.flatten(x, 1)\n",
+    "        x = self.fc(x)\n",
+    "        x = self.sigmoid(x)\n",
+    "        return x\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        return self._forward_impl(x)\n",
+    "\n",
+    "def ComplexResNet18():\n",
+    "    return ComplexResNet(BasicBlock, [2, 2, 2, 2])\n",
+    "\n",
+    "# Create the model instance\n",
+    "model = ComplexResNet18()\n",
+    "print(model)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4bc1b5d",
+   "metadata": {},
+   "source": [
+    "### Complex focal Loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "61c29429",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ComplexFocalLoss(nn.Module):\n",
+    "    def __init__(self, alpha=1, gamma=2, reduction='mean'):\n",
+    "        super(ComplexFocalLoss, self).__init__()\n",
+    "        self.alpha = alpha\n",
+    "        self.gamma = gamma\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "        \n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction='none')\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction='none')\n",
+    "        \n",
+    "        real_pt = torch.exp(-real_BCE_loss)\n",
+    "        imag_pt = torch.exp(-imag_BCE_loss)\n",
+    "        \n",
+    "        real_F_loss = self.alpha * (1 - real_pt) ** self.gamma * real_BCE_loss\n",
+    "        imag_F_loss = self.alpha * (1 - imag_pt) ** self.gamma * imag_BCE_loss\n",
+    "\n",
+    "        if self.reduction == 'mean':\n",
+    "            return (torch.mean(real_F_loss) + torch.mean(imag_F_loss)) / 2\n",
+    "        elif self.reduction == 'sum':\n",
+    "            return torch.sum(real_F_loss) + torch.sum(imag_F_loss)\n",
+    "        else:\n",
+    "            return real_F_loss + imag_F_loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abb35ba2",
+   "metadata": {},
+   "source": [
+    "### Training with complex focal loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86d7526b",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Initialize and train the ResNet-18 model\n",
+    "model = ComplexResNet18().to(device)\n",
+    "criterion = ComplexFocalLoss()\n",
+    "\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations =train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3)\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd0c9d58",
+   "metadata": {},
+   "source": [
+    "### CVNN RV-BCE and CV-BCE Loss function implementation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99c736b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# CV BCE Loss Function Definition\n",
+    "class ComplexValuedBCELoss(nn.Module):\n",
+    "    def __init__(self, reduction='mean'):\n",
+    "        super(ComplexValuedBCELoss, self).__init__()\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "\n",
+    "        # Calculate binary cross-entropy for both real and imaginary parts\n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction=self.reduction)\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction=self.reduction)\n",
+    "        \n",
+    "        # Combine the losses (you can adjust the weighting if necessary)\n",
+    "        combined_BCE_loss = (real_BCE_loss + imag_BCE_loss) / 2\n",
+    "        return combined_BCE_loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93d19ea7",
+   "metadata": {},
+   "source": [
+    "### CV-BCE Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c56d5b4",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Set the criterion for CV BCE\n",
+    "criterion = ComplexValuedBCELoss()\n",
+    "\n",
+    "# Train the ResNet-18 model with CV BCE\n",
+    "device = torch.device('cuda')\n",
+    "model = ComplexResNet18().to(device)\n",
+    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
+    "\n",
+    "# Start training with the previously defined train_model function\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations = train_model(\n",
+    "    model, train_loader, valid_loader, criterion, \n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3\n",
+    ")\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ccd50ff",
+   "metadata": {},
+   "source": [
+    "### Save and Plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb41b92f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Define save directory\n",
+    "save_dir = 'CMuSeNet_results/segmentation_OTA'\n",
+    "\n",
+    "# Create the directory if it doesn't exist\n",
+    "os.makedirs(save_dir, exist_ok=True)\n",
+    "\n",
+    "# Plot training loss\n",
+    "plt.figure()\n",
+    "plt.plot(range(1, len(train_losses) + 1), train_losses, label='Training Loss', color='blue')\n",
+    "plt.title('Training Loss')\n",
+    "plt.xlabel('Epoch')\n",
+    "plt.ylabel('Loss')\n",
+    "plt.legend()\n",
+    "\n",
+    "# Save the training loss figure as PNG and SVG\n",
+    "plt.savefig(os.path.join(save_dir, 'training_loss.png'))\n",
+    "plt.savefig(os.path.join(save_dir, 'training_loss.svg'))\n",
+    "\n",
+    "# Show the training loss plot\n",
+    "plt.show()\n",
+    "\n",
+    "# Plot validation accuracy\n",
+    "plt.figure()\n",
+    "plt.plot(range(1, len(val_accuracies) + 1), val_accuracies, label='Validation Accuracy', color='green')\n",
+    "plt.title('Validation Accuracy')\n",
+    "plt.xlabel('Epoch')\n",
+    "plt.ylabel('Accuracy')\n",
+    "plt.legend()\n",
+    "\n",
+    "# Save the validation accuracy figure as PNG and SVG\n",
+    "plt.savefig(os.path.join(save_dir, 'validation_accuracy.png'))\n",
+    "plt.savefig(os.path.join(save_dir, 'validation_accuracy.svg'))\n",
+    "\n",
+    "# Show the validation accuracy plot\n",
+    "plt.show()\n",
+    "\n",
+    "# Save the actual data to a JSON file\n",
+    "results = {\n",
+    "    \"train_losses\": train_losses,\n",
+    "    \"val_accuracies\": val_accuracies,\n",
+    "    \"epoch_durations\": epoch_durations,\n",
+    "    \"combined_epoch_time\": combined_epoch_time\n",
+    "}\n",
+    "\n",
+    "# Save JSON file\n",
+    "with open(os.path.join(save_dir, 'training_validation_results.json'), 'w') as f:\n",
+    "    json.dump(results, f)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a757949",
+   "metadata": {},
+   "source": [
+    "### Transfer Learning from Synthetic model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ee265d28",
+   "metadata": {},
+   "source": [
+    "### Load pre-trained model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0dec6746",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Block to load pre-trained model and prepare for transfer learning\n",
+    "device = torch.device(\"cuda\")\n",
+    "\n",
+    "# Load the pre-trained model\n",
+    "\n",
+    "model_path = \"path/to/model/save.pth\"\n",
+    "model = ComplexResNet18().to(device)\n",
+    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
+    "\n",
+    "# Freeze all layers except the final layer\n",
+    "for param in model.parameters():\n",
+    "    param.requires_grad = False\n",
+    "\n",
+    "# Modify the final layer for transfer learning (adjust `num_classes` as needed)\n",
+    "num_classes = STFT_LENGTH  # Set based on your current task\n",
+    "model.fc = cplx.ComplexLinear(512 * BasicBlock.expansion, num_classes).to(device)\n",
+    "\n",
+    "# Unfreeze the final layer for training\n",
+    "for param in model.fc.parameters():\n",
+    "    param.requires_grad = True\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "21e1e62b",
+   "metadata": {},
+   "source": [
+    "### Complex Learning for Transfer Learning (Same as above but easier access)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c6656d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ComplexFocalLoss(nn.Module):\n",
+    "    def __init__(self, alpha=0.5, gamma=2, reduction='mean'):\n",
+    "        super(ComplexFocalLoss, self).__init__()\n",
+    "        self.alpha = alpha\n",
+    "        self.gamma = gamma\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "        \n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction='none')\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction='none')\n",
+    "        \n",
+    "        real_pt = torch.exp(-real_BCE_loss)\n",
+    "        imag_pt = torch.exp(-imag_BCE_loss)\n",
+    "        \n",
+    "        real_F_loss = self.alpha * (1 - real_pt) ** self.gamma * real_BCE_loss\n",
+    "        imag_F_loss = self.alpha * (1 - imag_pt) ** self.gamma * imag_BCE_loss\n",
+    "\n",
+    "        if self.reduction == 'mean':\n",
+    "            return (torch.mean(real_F_loss) + torch.mean(imag_F_loss)) / 2\n",
+    "        elif self.reduction == 'sum':\n",
+    "            return torch.sum(real_F_loss) + torch.sum(imag_F_loss)\n",
+    "        else:\n",
+    "            return real_F_loss + imag_F_loss\n",
+    "\n",
+    "# Update the IoU calculation to handle complex values\n",
+    "def calculate_iou(pred, target, threshold=0.5):\n",
+    "    real_pred = (pred.real > threshold).float()\n",
+    "    imag_pred = (pred.imag > threshold).float()\n",
+    "    \n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred).float()\n",
+    "    \n",
+    "    intersection = (combined_pred * target).sum(dim=1)\n",
+    "    union = (combined_pred + target).sum(dim=1) - intersection\n",
+    "    iou = (intersection / union).mean().item()\n",
+    "    return iou"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bc9b7701",
+   "metadata": {},
+   "source": [
+    "### Transfer Learning"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c291a42e",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Define a new criterion and optimizer for fine-tuning\n",
+    "# You may select between Focal Loss or BCE as your criterion\n",
+    "#criterion = ComplexValuedBCELoss()  # or ComplexValuedBCELoss()\n",
+    "criterion = ComplexFocalLoss()\n",
+    "# Use a smaller learning rate for fine-tuning\n",
+    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)\n",
+    "\n",
+    "# Train the model (fine-tuning)\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations= train_model(\n",
+    "    model, train_loader, valid_loader, criterion,\n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3\n",
+    ")\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "98f81acc",
+   "metadata": {},
+   "source": [
+    "## Transfer Transfer Learning (Different Radio)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "55017794",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Block to load pre-trained model and prepare for transfer learning\n",
+    "device = torch.device(\"cuda\")\n",
+    "\n",
+    "model_path = \"/path/to/model/save.pth\"\n",
+    "model = ComplexResNet18().to(device)\n",
+    "#model = ComplexValuedBCELoss().to(device)\n",
+    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
+    "\n",
+    "# Freeze all layers except the final layer\n",
+    "for param in model.parameters():\n",
+    "    param.requires_grad = False\n",
+    "\n",
+    "# Modify the final layer for transfer learning (adjust `num_classes` as needed)\n",
+    "num_classes = STFT_LENGTH  # Set based on your current task\n",
+    "model.fc = cplx.ComplexLinear(512 * BasicBlock.expansion, num_classes).to(device)\n",
+    "\n",
+    "# Unfreeze the final layer for training\n",
+    "for param in model.fc.parameters():\n",
+    "    param.requires_grad = True\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5933b01f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define a new criterion and optimizer for fine-tuning\n",
+    "# You may select between Focal Loss or BCE as your criterion\n",
+    "#criterion = ComplexValuedBCELoss()  # or ComplexValuedBCELoss()\n",
+    "criterion = ComplexFocalLoss()\n",
+    "# Use a smaller learning rate for fine-tuning\n",
+    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)\n",
+    "\n",
+    "# Train the model (fine-tuning)\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations= train_model(\n",
+    "    model, train_loader, valid_loader, criterion,\n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=1\n",
+    ")\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dbef8bad",
+   "metadata": {},
+   "source": [
+    "### Evaluation CVNN OTA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0ac03b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from tqdm import tqdm\n",
+    "import numpy as np\n",
+    "from collections import defaultdict\n",
+    "import torch.nn.functional as F\n",
+    "from scipy.optimize import linear_sum_assignment\n",
+    "from torch.utils.data import ConcatDataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f831e874",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = \"cuda\"\n",
+    "\n",
+    "model_path = \"/path/to/model/save.pth\"\n",
+    "model = ComplexResNet18().to(device)\n",
+    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
+    "model.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a303080e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the pre-trained model for evaluation\n",
+    "\n",
+    "full_dataset = ConcatDataset([\n",
+    "    WidebandSignalDataset(signal_ids=train, return_snrs=True),\n",
+    "    WidebandSignalDataset(signal_ids=validation, return_snrs=True),\n",
+    "    WidebandSignalDataset(signal_ids=test, return_snrs=True)\n",
+    "])\n",
+    "full_loader = DataLoader(full_dataset, batch_size=64, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ad326f1d",
+   "metadata": {},
+   "source": [
+    "### Function initialization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00d0228c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def expand_true(array, distance=1):\n",
+    "    # Create kernel of appropriate size\n",
+    "    kernel = torch.ones((1, 1, distance * 2 + 1), device=array.device)\n",
+    "    array = array.unsqueeze(1).float()  # Add channel dimension\n",
+    "    result = F.conv1d(array, kernel, padding=distance)\n",
+    "    result = result.squeeze(1)  # Remove the extra dimension\n",
+    "    return result > 0\n",
+    "def reshape_to_2d(data):\n",
+    "    return data.view(-1, 1, 128, 128)  # Reshape to [batch, channels, height, width]\n",
+    "def get_true_groups(tensor, device):\n",
+    "    assert tensor.dim() == 2, 'This function handles 2D tensor only'\n",
+    "    all_groups = []\n",
+    "    for i in range(tensor.size(0)):\n",
+    "        item = tensor[i]\n",
+    "        item = torch.cat([torch.tensor([False]).to(device), item, torch.tensor([False]).to(device)])\n",
+    "        diffs = item.float().diff()\n",
+    "        starts = (diffs == 1).nonzero(as_tuple=True)[0]\n",
+    "        ends = (diffs == -1).nonzero(as_tuple=True)[0] - 1\n",
+    "        groups = [(start.item(), end.item()) for start, end in zip(starts, ends)]\n",
+    "        all_groups.append(groups)\n",
+    "    return all_groups\n",
+    "\n",
+    "def calculate_iou(box1, box2):\n",
+    "    intersection = max(0, min(box1[1], box2[1]) - max(box1[0], box2[0]))\n",
+    "    union = max(box1[1], box2[1]) - min(box1[0], box2[0])\n",
+    "    return intersection / union if union != 0 else 0\n",
+    "\n",
+    "def match_targets(targets, preds):\n",
+    "    ious = []\n",
+    "    for target in targets:\n",
+    "        iou_targets = []\n",
+    "        for pred in preds:\n",
+    "            iou_targets.append(calculate_iou(target, pred))\n",
+    "        ious.append(iou_targets)\n",
+    "    cost_matrix = np.array(ious)\n",
+    "    row_ind, col_ind = linear_sum_assignment(-cost_matrix)\n",
+    "    return row_ind, col_ind\n",
+    "\n",
+    "def calculate_matched_ious(target_boxes, prediction_boxes, matching):\n",
+    "    ious = [0 for _ in target_boxes]\n",
+    "    matching_dict = dict(zip(*matching))\n",
+    "    for target_index, target_box in enumerate(target_boxes):\n",
+    "        if target_index in matching_dict:\n",
+    "            pred_index = matching_dict[target_index]\n",
+    "            if pred_index < len(prediction_boxes):\n",
+    "                box1 = target_box\n",
+    "                box2 = prediction_boxes[pred_index]\n",
+    "                ious[target_index] = calculate_iou(box1, box2)\n",
+    "    return ious\n",
+    "def model_predictor(signals):\n",
+    "    # Convert signals to complex tensors\n",
+    "    if signals.dtype != torch.complex64 and signals.dtype != torch.complex128:\n",
+    "        signals = signals.type(torch.complex64)\n",
+    "    # Reshape the input signals to the expected shape\n",
+    "    signals = reshape_to_2d(signals)\n",
+    "    signals = signals.to(device)\n",
+    "    # Use the already loaded model and apply thresholding\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model(signals)\n",
+    "    # Handle complex outputs appropriately\n",
+    "    real_outputs = outputs.real\n",
+    "    imag_outputs = outputs.imag\n",
+    "    real_pred = (real_outputs > 0.5)\n",
+    "    imag_pred = (imag_outputs > 0.5)\n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred)\n",
+    "    return expand_true(combined_pred.float())\n",
+    "\n",
+    "# Complex IoU Implementation\n",
+    "def calculate_complex_iou(box1_real, box1_imag, box2_real, box2_imag):\n",
+    "    # Calculate real component intersection\n",
+    "    real_intersection = max(0, min(box1_real[1], box2_real[1]) - max(box1_real[0], box2_real[0]))\n",
+    "    real_union = max(box1_real[1], box2_real[1]) - min(box1_real[0], box2_real[0])\n",
+    "    \n",
+    "    # Calculate imaginary component intersection\n",
+    "    imag_intersection = max(0, min(box1_imag[1], box2_imag[1]) - max(box1_imag[0], box2_imag[0]))\n",
+    "    imag_union = max(box1_imag[1], box2_imag[1]) - min(box1_imag[0], box2_imag[0])\n",
+    "    \n",
+    "    # Combine intersections and unions\n",
+    "    total_intersection = real_intersection + imag_intersection\n",
+    "    total_union = real_union + imag_union\n",
+    "    \n",
+    "    # Return IoU\n",
+    "    return total_intersection / total_union if total_union != 0 else 0\n",
+    "\n",
+    "def match_complex_targets(targets_real, targets_imag, preds_real, preds_imag):\n",
+    "    ious = []\n",
+    "    for target_real, target_imag in zip(targets_real, targets_imag):\n",
+    "        iou_targets = []\n",
+    "        for pred_real, pred_imag in zip(preds_real, preds_imag):\n",
+    "            iou_targets.append(calculate_complex_iou(target_real, target_imag, pred_real, pred_imag))\n",
+    "        ious.append(iou_targets)\n",
+    "    cost_matrix = np.array(ious)\n",
+    "    row_ind, col_ind = linear_sum_assignment(-cost_matrix)\n",
+    "    return row_ind, col_ind\n",
+    "\n",
+    "def calculate_matched_complex_ious(target_boxes_real, target_boxes_imag, \n",
+    "                                   prediction_boxes_real, prediction_boxes_imag, matching):\n",
+    "    ious = [0 for _ in target_boxes_real]\n",
+    "    matching_dict = dict(zip(*matching))\n",
+    "    for target_index, (target_box_real, target_box_imag) in enumerate(zip(target_boxes_real, target_boxes_imag)):\n",
+    "        if target_index in matching_dict:\n",
+    "            pred_index = matching_dict[target_index]\n",
+    "            if pred_index < len(prediction_boxes_real):\n",
+    "                box1_real, box1_imag = target_box_real, target_box_imag\n",
+    "                box2_real, box2_imag = prediction_boxes_real[pred_index], prediction_boxes_imag[pred_index]\n",
+    "                ious[target_index] = calculate_complex_iou(box1_real, box1_imag, box2_real, box2_imag)\n",
+    "    return ious\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c114c7a2",
+   "metadata": {},
+   "source": [
+    "### Evaluate function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "41f12e83",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def evaluate(predictor, data_loader, device=\"cuda\"):\n",
+    "    iou_thresholds = [0.5, 0.7, 0.9]\n",
+    "    snr_metrics = defaultdict(lambda: {\n",
+    "        \"iou_sum\": 0.0,\n",
+    "        \"iou_count\": 0,\n",
+    "        \"recall_counts\": defaultdict(int),\n",
+    "        \"total_samples\": defaultdict(int),\n",
+    "        \"correct_pixels\": 0,\n",
+    "        \"total_pixels\": 0\n",
+    "    })\n",
+    "    total_iou_sum, total_iou_count = 0.0, 0\n",
+    "    total_correct_pixels, total_total_pixels = 0, 0\n",
+    "    total_recall_counts = defaultdict(int)\n",
+    "    total_samples = defaultdict(int)\n",
+    "\n",
+    "    for batch in tqdm(data_loader, desc=\"Evaluating\"):\n",
+    "        if len(batch) == 3:\n",
+    "            inputs, masks, snrs_in_batch = batch\n",
+    "        else:\n",
+    "            inputs, masks = batch\n",
+    "            snrs_in_batch = [0] * len(inputs)  # Default SNR if not provided\n",
+    "\n",
+    "        inputs = inputs.to(device)\n",
+    "        masks = masks.to(device)\n",
+    "        outputs = predictor(inputs)\n",
+    "\n",
+    "        for i in range(len(inputs)):\n",
+    "            mask = masks[i]\n",
+    "            output = outputs[i]\n",
+    "\n",
+    "            # Resize output to match mask shape if necessary\n",
+    "            if output.numel() != mask.numel():\n",
+    "                output = output.expand_as(mask) if output.numel() == 1 else output.reshape_as(mask)\n",
+    "\n",
+    "            thresholded_output = (output >= 0.5).float()\n",
+    "\n",
+    "            correct_pixels = (thresholded_output == mask).sum().item()\n",
+    "            total_pixels = mask.numel()\n",
+    "            total_correct_pixels += correct_pixels\n",
+    "            total_total_pixels += total_pixels\n",
+    "\n",
+    "            # Get SNR value and round it to the nearest integer\n",
+    "            snr = snrs_in_batch[i]\n",
+    "            if isinstance(snr, torch.Tensor):\n",
+    "                snr = snr.item()\n",
+    "            snr = int(round(snr))  # Round SNR to the nearest integer\n",
+    "\n",
+    "            snr_metrics[snr][\"correct_pixels\"] += correct_pixels\n",
+    "            snr_metrics[snr][\"total_pixels\"] += total_pixels\n",
+    "\n",
+    "            target_boxes = get_true_groups(mask.unsqueeze(0), device=device)[0]\n",
+    "            pred_boxes = get_true_groups(thresholded_output.unsqueeze(0), device=device)[0]\n",
+    "            if not target_boxes or not pred_boxes:\n",
+    "                continue\n",
+    "            matching = match_targets(target_boxes, pred_boxes)\n",
+    "            matched_ious = calculate_matched_ious(target_boxes, pred_boxes, matching)\n",
+    "\n",
+    "            snr_metrics[snr][\"iou_sum\"] += sum(matched_ious)\n",
+    "            snr_metrics[snr][\"iou_count\"] += len(matched_ious)\n",
+    "            total_iou_sum += sum(matched_ious)\n",
+    "            total_iou_count += len(matched_ious)\n",
+    "\n",
+    "            for th in iou_thresholds:\n",
+    "                true_positives = sum(1 for iou in matched_ious if iou >= th)\n",
+    "                snr_metrics[snr][\"recall_counts\"][th] += true_positives\n",
+    "                snr_metrics[snr][\"total_samples\"][th] += len(target_boxes)\n",
+    "                total_recall_counts[th] += true_positives\n",
+    "                total_samples[th] += len(target_boxes)\n",
+    "\n",
+    "    # Calculate overall metrics\n",
+    "    overall_accuracy = (total_correct_pixels / total_total_pixels) * 100 if total_total_pixels > 0 else 0\n",
+    "    overall_iou = total_iou_sum / total_iou_count if total_iou_count > 0 else 0\n",
+    "    overall_recall = {\n",
+    "        th: total_recall_counts[th] / total_samples[th] if total_samples[th] > 0 else 0\n",
+    "        for th in iou_thresholds\n",
+    "    }\n",
+    "\n",
+    "    # Print overall results\n",
+    "    print(f\"Overall Accuracy: {overall_accuracy:.2f}%\")\n",
+    "    print(f\"Overall IoU Score: {overall_iou:.4f}\")\n",
+    "    for th in iou_thresholds:\n",
+    "        print(f\"Recall at threshold {th}: {overall_recall[th]:.4f}\")\n",
+    "\n",
+    "    # Print per-SNR results\n",
+    "    for snr in sorted(snr_metrics.keys()):\n",
+    "        metrics = snr_metrics[snr]\n",
+    "        snr_accuracy = (metrics[\"correct_pixels\"] / metrics[\"total_pixels\"]) * 100 if metrics[\"total_pixels\"] > 0 else 0\n",
+    "        snr_iou = metrics[\"iou_sum\"] / metrics[\"iou_count\"] if metrics[\"iou_count\"] > 0 else 0\n",
+    "        print(f\"SNR: {snr} dB - Accuracy: {snr_accuracy:.2f}%\")\n",
+    "        print(f\"   IoU: {snr_iou:.4f}\")\n",
+    "        for th in iou_thresholds:\n",
+    "            recall = metrics[\"recall_counts\"][th] / metrics[\"total_samples\"][th] if metrics[\"total_samples\"][th] > 0 else 0\n",
+    "            print(f\"   Recall at threshold {th}: {recall:.4f}\")\n",
+    "\n",
+    "    return snr_metrics\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d2fd13f",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "# Run evaluation on the full dataset\n",
+    "snr_metrics = evaluate(model_predictor, full_loader, device=device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "07eade04",
+   "metadata": {},
+   "source": [
+    "### Save and Plot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bc84b73a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "def save_results_and_plot(snr_metrics, save_path):\n",
+    "    \"\"\"\n",
+    "    Saves evaluation results to a JSON file and generates plots for Accuracy, IoU, and Recall vs. SNR.\n",
+    "    Sets x-axis limits to range from -9 dB to 12 dB to eliminate blank space on the right.\n",
+    "\n",
+    "    Args:\n",
+    "        snr_metrics (dict): The evaluation results obtained from the evaluate function.\n",
+    "        save_path (str): The directory path where results and plots will be saved.\n",
+    "\n",
+    "    Outputs:\n",
+    "        - evaluation_results.json\n",
+    "        - accuracy_vs_snr.png and .svg\n",
+    "        - iou_vs_snr.png and .svg\n",
+    "        - recall_vs_snr.png and .svg\n",
+    "    \"\"\"\n",
+    "    # Ensure the directory exists\n",
+    "    os.makedirs(save_path, exist_ok=True)\n",
+    "    \n",
+    "    # Extract data from snr_metrics\n",
+    "    snr_list = sorted(snr_metrics.keys())\n",
+    "    accuracy_list = []\n",
+    "    iou_list = []\n",
+    "    recall_05 = []\n",
+    "    recall_07 = []\n",
+    "    recall_09 = []\n",
+    "    \n",
+    "    # Prepare data for JSON serialization\n",
+    "    json_data = {}\n",
+    "    \n",
+    "    for snr in snr_list:\n",
+    "        metrics = snr_metrics[snr]\n",
+    "        snr_accuracy = (metrics[\"correct_pixels\"] / metrics[\"total_pixels\"]) * 100 if metrics[\"total_pixels\"] > 0 else 0\n",
+    "        snr_iou = metrics[\"iou_sum\"] / metrics[\"iou_count\"] if metrics[\"iou_count\"] > 0 else 0\n",
+    "        recall_at_05 = metrics[\"recall_counts\"][0.5] / metrics[\"total_samples\"][0.5] if metrics[\"total_samples\"][0.5] > 0 else 0\n",
+    "        recall_at_07 = metrics[\"recall_counts\"][0.7] / metrics[\"total_samples\"][0.7] if metrics[\"total_samples\"][0.7] > 0 else 0\n",
+    "        recall_at_09 = metrics[\"recall_counts\"][0.9] / metrics[\"total_samples\"][0.9] if metrics[\"total_samples\"][0.9] > 0 else 0\n",
+    "\n",
+    "        # Append to lists for plotting\n",
+    "        accuracy_list.append(snr_accuracy)\n",
+    "        iou_list.append(snr_iou)\n",
+    "        recall_05.append(recall_at_05)\n",
+    "        recall_07.append(recall_at_07)\n",
+    "        recall_09.append(recall_at_09)\n",
+    "\n",
+    "        # Prepare data for JSON\n",
+    "        json_data[snr] = {\n",
+    "            \"accuracy\": snr_accuracy,\n",
+    "            \"iou\": snr_iou,\n",
+    "            \"recall\": {\n",
+    "                \"0.5\": recall_at_05,\n",
+    "                \"0.7\": recall_at_07,\n",
+    "                \"0.9\": recall_at_09,\n",
+    "            }\n",
+    "        }\n",
+    "    \n",
+    "    # Save json_data to JSON file\n",
+    "    json_file_path = os.path.join(save_path, 'evaluation_results.json')\n",
+    "    with open(json_file_path, 'w') as json_file:\n",
+    "        json.dump(json_data, json_file, indent=4)\n",
+    "    \n",
+    "    # Plot Accuracy vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, accuracy_list, marker='o', label='Accuracy')\n",
+    "    plt.title('Accuracy vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('Accuracy (%)')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    #plt.xlim(-9, 12)\n",
+    "    plt.xlim(-16, 16)\n",
+    "    # Save the plot\n",
+    "    accuracy_png_path = os.path.join(save_path, 'accuracy_vs_snr.png')\n",
+    "    accuracy_svg_path = os.path.join(save_path, 'accuracy_vs_snr.svg')\n",
+    "    plt.savefig(accuracy_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(accuracy_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n",
+    "    \n",
+    "    # Plot IoU vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, iou_list, marker='o', color='orange', label='IoU')\n",
+    "    plt.title('IoU vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('IoU')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    #plt.xlim(-9, 12)\n",
+    "    plt.xlim(-16, 16)\n",
+    "    # Save the plot\n",
+    "    iou_png_path = os.path.join(save_path, 'iou_vs_snr.png')\n",
+    "    iou_svg_path = os.path.join(save_path, 'iou_vs_snr.svg')\n",
+    "    plt.savefig(iou_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(iou_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n",
+    "    \n",
+    "    # Plot Recall at Different IoU Thresholds vs. SNR\n",
+    "    plt.figure(figsize=(10, 6))\n",
+    "    plt.plot(snr_list, recall_05, marker='o', label='Recall @ IoU 0.5')\n",
+    "    plt.plot(snr_list, recall_07, marker='s', label='Recall @ IoU 0.7')\n",
+    "    plt.plot(snr_list, recall_09, marker='^', label='Recall @ IoU 0.9')\n",
+    "    plt.title('Recall at Different IoU Thresholds vs. SNR')\n",
+    "    plt.xlabel('SNR (dB)')\n",
+    "    plt.ylabel('Recall')\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    # Set x-axis limits\n",
+    "    plt.xlim(-9, 12)\n",
+    "    \n",
+    "    # Save the plot\n",
+    "    recall_png_path = os.path.join(save_path, 'recall_vs_snr.png')\n",
+    "    recall_svg_path = os.path.join(save_path, 'recall_vs_snr.svg')\n",
+    "    plt.savefig(recall_png_path, format='png', bbox_inches='tight')\n",
+    "    plt.savefig(recall_svg_path, format='svg', bbox_inches='tight')\n",
+    "    \n",
+    "    plt.show()\n",
+    "    plt.close()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1974e70d",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "save_path = 'CMuSeNet_results/OTA'\n",
+    "\n",
+    "# Save results and generate plots\n",
+    "save_results_and_plot(snr_metrics, save_path)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

CMuSeNet_Synthetic.ipynb

@@ -0,0 +1,1241 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b5007b71",
+   "metadata": {},
+   "source": [
+    "### Initialization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3e6b1226",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Initialization block\n",
+    "from pathlib import Path\n",
+    "import numpy as np\n",
+    "import json\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "import math\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
+    "\n",
+    "STFT_LENGTH = 16 * 1024\n",
+    "DATA_DIR = Path(\"dataset/\")\n",
+    "SAMPLE_RATE = 20e6\n",
+    "MODULATIONS = [\"QPSK\", \"BPSK\", \"8-PSK\", \"8-QAM\", \"16-QAM\", \"GMSK\", \"2-FSK\"]\n",
+    "MODULATION_LABELS = {j: i for i, j in enumerate(MODULATIONS)}\n",
+    "NUMBER_OF_MODULATIONS = len(MODULATIONS)\n",
+    "\n",
+    "def load_data(snr, name, load_metadata_only=False):\n",
+    "    if not load_metadata_only:\n",
+    "        with open(DATA_DIR/str(snr)/str(name)/\"data.dat\", \"rb\") as f:\n",
+    "            signal = np.fromfile(f, dtype=np.complex128)\n",
+    "    else:\n",
+    "        signal = None\n",
+    "    with open(DATA_DIR/str(snr)/str(name)/\"meta-data.json\") as f:\n",
+    "        meta = json.load(f)\n",
+    "        if type(meta) == dict:\n",
+    "            meta = [meta]\n",
+    "    return signal, meta\n",
+    "\n",
+    "    \n",
+    "def _get_all_numbered_dirs(root_dir):\n",
+    "    dirs = []\n",
+    "    for directory in root_dir.iterdir():\n",
+    "        dirs.append(int(directory.name))\n",
+    "    dirs.sort()\n",
+    "    return dirs\n",
+    "\n",
+    "def get_signals(snr):\n",
+    "    return _get_all_numbered_dirs(Path(DATA_DIR)/str(snr))\n",
+    "\n",
+    "\n",
+    "def get_snrs(root_dir=DATA_DIR):\n",
+    "    return _get_all_numbered_dirs(root_dir)\n",
+    "        \n",
+    "        \n",
+    "def process_metadata(metadata):\n",
+    "    scaled_metadata =  [\n",
+    "        {\n",
+    "            \"position\": (SAMPLE_RATE/2 + i['fc'], i['bw']),\n",
+    "            \"mod\": i[\"mod\"]\n",
+    "        }\n",
+    "        for i in metadata\n",
+    "    ]\n",
+    "    return scaled_metadata\n",
+    "\n",
+    "\n",
+    "def process_signal(signal):\n",
+    "    signal = signal[:STFT_LENGTH]\n",
+    "\n",
+    "    signal = np.fft.fft(signal)\n",
+    "    signal = np.fft.fftshift(signal)\n",
+    "    signal /= np.max(np.abs(signal))\n",
+    "    \n",
+    "    #return np.expand_dims(signal, axis=0)\n",
+    "    return signal"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "440b802c",
+   "metadata": {},
+   "source": [
+    "### Data Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31bc3770",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MASK_SIZE = int(STFT_LENGTH)\n",
+    "\n",
+    "class WidebandSignalDataset(torch.utils.data.Dataset):\n",
+    "    def __init__(self, signal_ids, mask_size=MASK_SIZE, return_snr=False):\n",
+    "        self.mask_size = mask_size\n",
+    "        self.signal_ids = signal_ids\n",
+    "        self.return_snr = return_snr  # New parameter to control SNR return\n",
+    "        loaded_data = []\n",
+    "        for snr, signal_id in tqdm(self.signal_ids):\n",
+    "            signal, masks = self.process_signal(snr, signal_id)\n",
+    "            loaded_data.append((signal, masks))\n",
+    "        self.loaded_data = loaded_data\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.signal_ids)\n",
+    "\n",
+    "    def __getitem__(self, index):\n",
+    "        signal, masks = self.loaded_data[index]\n",
+    "        if self.return_snr:\n",
+    "            snr, _ = self.signal_ids[index]\n",
+    "            return signal, masks, snr  # Return SNR during evaluation\n",
+    "        else:\n",
+    "            return signal, masks  # Return only signal and masks during training\n",
+    "\n",
+    "    def process_signal(self, snr, signal_id):\n",
+    "        signal, metadata = load_data(snr, signal_id)\n",
+    "        scaled_metadata = process_metadata(metadata)\n",
+    "        signal = process_signal(signal)\n",
+    "        signal = torch.from_numpy(signal)\n",
+    "        masks = torch.zeros(self.mask_size)\n",
+    "        scale_ratio = self.mask_size / SAMPLE_RATE\n",
+    "        for meta in scaled_metadata:\n",
+    "            f, b = meta['position']\n",
+    "            x1, x2 = math.floor((f - b / 2) * scale_ratio), math.ceil((f + b / 2) * scale_ratio)\n",
+    "            masks[x1:x2] = 1\n",
+    "        return signal.type(torch.complex64), masks.type(torch.FloatTensor)\n",
+    "\n",
+    "# Train test split 80 - 10 - 10\n",
+    "train, test, validation = [], [], [] \n",
+    "for snr in get_snrs():\n",
+    "    signals = get_signals(snr)\n",
+    "    total_signals = len(signals)\n",
+    "    for signal in signals:\n",
+    "        if signal <= 0.8 * total_signals:\n",
+    "            train.append((snr, signal))\n",
+    "        elif signal <= 0.9 * total_signals:\n",
+    "            validation.append((snr, signal))\n",
+    "        else:\n",
+    "            test.append((snr, signal))\n",
+    "            \n",
+    "print(\"Train\", len(train))\n",
+    "print(\"Validation\", len(validation))\n",
+    "print(\"Test\", len(test))\n",
+    "\n",
+    "train_dataset = WidebandSignalDataset(signal_ids=train)\n",
+    "validation_dataset = WidebandSignalDataset(signal_ids=validation)\n",
+    "test_dataset = WidebandSignalDataset(signal_ids=test)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "637ae774",
+   "metadata": {},
+   "source": [
+    "### Batch Loading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9af2450",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 64  # Updated batch size\n",
+    "\n",
+    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
+    "valid_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=False)\n",
+    "\n",
+    "print(\"Train labels shape:\", len(train_dataset))\n",
+    "print(\"Validation labels shape:\", len(validation_dataset))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a8e09e4",
+   "metadata": {},
+   "source": [
+    "### Early Stop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24f79a24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "class EarlyStopping:\n",
+    "    def __init__(self, patience=10, verbose=False, delta=0.0001, save_path='./models/CMuSeNet'):\n",
+    "        self.patience = patience\n",
+    "        self.verbose = verbose\n",
+    "        self.delta = delta\n",
+    "        self.counter = 0\n",
+    "        self.best_score = None\n",
+    "        self.early_stop = False\n",
+    "        self.val_loss_min = float('inf')\n",
+    "        self.best_model = None\n",
+    "        self.save_path = save_path\n",
+    "        os.makedirs(save_path, exist_ok=True)\n",
+    "        \n",
+    "    def __call__(self, val_loss, model):\n",
+    "        score = -val_loss\n",
+    "\n",
+    "        if self.best_score is None:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "        elif score < self.best_score + self.delta:\n",
+    "            self.counter += 1\n",
+    "            if self.verbose:\n",
+    "                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n",
+    "            if self.counter >= self.patience:\n",
+    "                self.early_stop = True\n",
+    "        else:\n",
+    "            self.best_score = score\n",
+    "            self.save_checkpoint(val_loss, model)\n",
+    "            self.counter = 0\n",
+    "\n",
+    "    def save_checkpoint(self, val_loss, model):\n",
+    "        if self.verbose:\n",
+    "            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')\n",
+    "        self.val_loss_min = val_loss\n",
+    "        self.best_model = model.state_dict()\n",
+    "        save_path = os.path.join(self.save_path, 'best_model.pth')\n",
+    "        torch.save(self.best_model, save_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c3fda74",
+   "metadata": {},
+   "source": [
+    "### Reshape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fcf91db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch.nn as nn\n",
+    "import complexPyTorch.complexLayers as cplx\n",
+    "import torch.nn.functional as F\n",
+    "import torch\n",
+    "\n",
+    "def reshape_to_2d(data):\n",
+    "    return data.view(-1, 1, 128, 128)  # Reshape to [batch, channels, height, width]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7d7562c",
+   "metadata": {},
+   "source": [
+    "### Complex IoU"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7218c3f3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def calculate_iou(pred, target, threshold=0.5):\n",
+    "    real_pred = (pred.real > threshold).float()\n",
+    "    imag_pred = (pred.imag > threshold).float()\n",
+    "    \n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred).float()\n",
+    "    \n",
+    "    intersection = (combined_pred * target).sum(dim=1)\n",
+    "    union = (combined_pred + target).sum(dim=1) - intersection\n",
+    "    iou = (intersection / union).mean().item()\n",
+    "    return iou"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "64f4063c",
+   "metadata": {},
+   "source": [
+    "### Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "66825110",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "\n",
+    "def validate_model(model, valid_loader, criterion):\n",
+    "    model.eval()\n",
+    "    running_loss = 0.0\n",
+    "    iou_scores = []\n",
+    "    total_correct = 0\n",
+    "    total_samples = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, masks in tqdm(valid_loader, desc=\"Validating\"):\n",
+    "            inputs = reshape_to_2d(inputs).to(device)\n",
+    "            masks = masks.to(device)\n",
+    "            outputs = model(inputs)\n",
+    "            loss = criterion(outputs, masks)\n",
+    "            running_loss += loss.item()\n",
+    "\n",
+    "            # Calculate IoU\n",
+    "            iou = calculate_iou(outputs, masks, threshold=0.5)\n",
+    "            iou_scores.append(iou)\n",
+    "            \n",
+    "            # Calculate accuracy\n",
+    "            preds = ((outputs.real > 0.5) & (outputs.imag > 0.5)).float()\n",
+    "            correct = (preds == masks).float().sum()\n",
+    "            total_correct += correct.item()\n",
+    "            total_samples += masks.numel()\n",
+    "\n",
+    "    val_loss = running_loss / len(valid_loader)\n",
+    "    mean_iou = sum(iou_scores) / len(iou_scores)\n",
+    "    accuracy = total_correct / total_samples * 100\n",
+    "\n",
+    "    print(f'Validation Loss: {val_loss:.6f}')\n",
+    "    print(f'Validation Accuracy: {accuracy:.2f}%')\n",
+    "\n",
+    "    return val_loss, accuracy\n",
+    "\n",
+    "def train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.00001], num_epochs=50, patience=3):\n",
+    "    train_losses = []\n",
+    "    val_losses = []\n",
+    "    val_accuracies = []\n",
+    "    epoch_durations = []\n",
+    "    \n",
+    "    current_lr = initial_lr\n",
+    "    for lr in lr_steps:\n",
+    "        optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
+    "        early_stopping = EarlyStopping(patience=patience, verbose=True, delta=0.001)\n",
+    "        print(\"Current learning rate: \", lr)\n",
+    "        for epoch in range(num_epochs):\n",
+    "            epoch_start_time = time.time()\n",
+    "            \n",
+    "            model.train()\n",
+    "            running_loss = 0.0\n",
+    "            for inputs, masks in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs} - Training\"):\n",
+    "                inputs = reshape_to_2d(inputs).to(device)\n",
+    "                masks = masks.to(device)\n",
+    "                outputs = model(inputs)\n",
+    "                loss = criterion(outputs, masks)\n",
+    "\n",
+    "                optimizer.zero_grad()\n",
+    "                loss.backward()\n",
+    "                optimizer.step()\n",
+    "\n",
+    "                running_loss += loss.item()\n",
+    "\n",
+    "            epoch_loss = running_loss / len(train_loader)\n",
+    "            train_losses.append(epoch_loss)\n",
+    "            print(f\"Training Loss: {epoch_loss:.6f}\")\n",
+    "\n",
+    "            val_loss, val_accuracy = validate_model(model, valid_loader, criterion)\n",
+    "            val_losses.append(val_loss)\n",
+    "            val_accuracies.append(val_accuracy)\n",
+    "            early_stopping(val_loss, model)\n",
+    "\n",
+    "            if early_stopping.early_stop:\n",
+    "                print(\"Early stopping triggered\")\n",
+    "                break\n",
+    "\n",
+    "            epoch_duration = time.time() - epoch_start_time\n",
+    "            epoch_durations.append(epoch_duration)\n",
+    "        if early_stopping.best_model is not None:\n",
+    "            print(f\"Loading best model from lr {lr}\")\n",
+    "            model.load_state_dict(early_stopping.best_model)\n",
+    "        \n",
+    "    print(\"Training completed.\")\n",
+    "    print(\"Epoch durations:\", epoch_durations)\n",
+    "    return model, train_losses, val_losses, val_accuracies, epoch_durations"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b80cb51",
+   "metadata": {},
+   "source": [
+    "### ResNet-18"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d208cb9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import complexPyTorch.complexLayers as cplx\n",
+    "from typing import Optional, Callable, Type, Union, List\n",
+    "import torch.nn.functional as F\n",
+    "from torch import Tensor\n",
+    "\n",
+    "def conv3x3(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"3x3 convolution with padding\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
+    "\n",
+    "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> cplx.ComplexConv2d:\n",
+    "    \"\"\"1x1 convolution\"\"\"\n",
+    "    return cplx.ComplexConv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
+    "\n",
+    "class BasicBlock(nn.Module):\n",
+    "    expansion = 1\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(BasicBlock, self).__init__()\n",
+    "        self.conv1 = conv3x3(inplanes, planes, stride)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.conv2 = conv3x3(planes, planes)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class Bottleneck(nn.Module):\n",
+    "    expansion = 4\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        inplanes: int,\n",
+    "        planes: int,\n",
+    "        stride: int = 1,\n",
+    "        downsample: Optional[nn.Module] = None,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(Bottleneck, self).__init__()\n",
+    "        self.conv1 = conv1x1(inplanes, planes)\n",
+    "        self.bn1 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv2 = conv3x3(planes, planes, stride)\n",
+    "        self.bn2 = cplx.ComplexBatchNorm2d(planes)\n",
+    "        self.conv3 = conv1x1(planes, planes * self.expansion)\n",
+    "        self.bn3 = cplx.ComplexBatchNorm2d(planes * self.expansion)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.downsample = downsample\n",
+    "        self.stride = stride\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        identity = x\n",
+    "\n",
+    "        out = self.conv1(x)\n",
+    "        out = self.bn1(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv2(out)\n",
+    "        out = self.bn2(out)\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        out = self.conv3(out)\n",
+    "        out = self.bn3(out)\n",
+    "\n",
+    "        if self.downsample is not None:\n",
+    "            identity = self.downsample(x)\n",
+    "\n",
+    "        out += identity\n",
+    "        out = self.relu(out)\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "class ComplexResNet(nn.Module):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        block: Type[Union[BasicBlock, Bottleneck]],\n",
+    "        layers: List[int],\n",
+    "        num_classes: int = STFT_LENGTH,\n",
+    "        zero_init_residual: bool = False,\n",
+    "        groups: int = 1,\n",
+    "        width_per_group: int = 64,\n",
+    "        norm_layer: Optional[Callable[..., nn.Module]] = None,\n",
+    "    ) -> None:\n",
+    "        super(ComplexResNet, self).__init__()\n",
+    "        if norm_layer is None:\n",
+    "            norm_layer = cplx.ComplexBatchNorm2d\n",
+    "        self._norm_layer = norm_layer\n",
+    "\n",
+    "        self.inplanes = 64\n",
+    "        self.dilation = 1\n",
+    "\n",
+    "        self.groups = groups\n",
+    "        self.base_width = width_per_group\n",
+    "        self.conv1 = cplx.ComplexConv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)\n",
+    "        self.bn1 = norm_layer(self.inplanes)\n",
+    "        self.relu = cplx.ComplexReLU()\n",
+    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
+    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
+    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
+    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
+    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
+    "        self.fc = cplx.ComplexLinear(512 * block.expansion, num_classes)\n",
+    "        self.sigmoid = cplx.ComplexSigmoid()\n",
+    "\n",
+    "    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int, stride: int = 1) -> nn.Sequential:\n",
+    "        norm_layer = self._norm_layer\n",
+    "        downsample = None\n",
+    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
+    "            downsample = nn.Sequential(\n",
+    "                conv1x1(self.inplanes, planes * block.expansion, stride),\n",
+    "                norm_layer(planes * block.expansion),\n",
+    "            )\n",
+    "\n",
+    "        layers = []\n",
+    "        layers.append(block(self.inplanes, planes, stride, downsample, norm_layer))\n",
+    "        self.inplanes = planes * block.expansion\n",
+    "        for _ in range(1, blocks):\n",
+    "            layers.append(block(self.inplanes, planes, norm_layer=norm_layer))\n",
+    "\n",
+    "        return nn.Sequential(*layers)\n",
+    "\n",
+    "    def _forward_impl(self, x: Tensor) -> Tensor:\n",
+    "        x = self.conv1(x)\n",
+    "        x = self.bn1(x)\n",
+    "        x = self.relu(x)\n",
+    "\n",
+    "        x = self.layer1(x)\n",
+    "        x = self.layer2(x)\n",
+    "        x = self.layer3(x)\n",
+    "        x = self.layer4(x)\n",
+    "\n",
+    "        x = self.avgpool(x)\n",
+    "        x = torch.flatten(x, 1)\n",
+    "        x = self.fc(x)\n",
+    "        x = self.sigmoid(x)\n",
+    "        return x\n",
+    "\n",
+    "    def forward(self, x: Tensor) -> Tensor:\n",
+    "        return self._forward_impl(x)\n",
+    "\n",
+    "def ComplexResNet18():\n",
+    "    return ComplexResNet(BasicBlock, [2, 2, 2, 2])\n",
+    "\n",
+    "# Create the model instance\n",
+    "model = ComplexResNet18()\n",
+    "print(model)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4bc1b5d",
+   "metadata": {},
+   "source": [
+    "### Complex focal Loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "61c29429",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class ComplexFocalLoss(nn.Module):\n",
+    "    def __init__(self, alpha=1, gamma=2, reduction='mean'):\n",
+    "        super(ComplexFocalLoss, self).__init__()\n",
+    "        self.alpha = alpha\n",
+    "        self.gamma = gamma\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "        \n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction='none')\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction='none')\n",
+    "        \n",
+    "        real_pt = torch.exp(-real_BCE_loss)\n",
+    "        imag_pt = torch.exp(-imag_BCE_loss)\n",
+    "        \n",
+    "        real_F_loss = self.alpha * (1 - real_pt) ** self.gamma * real_BCE_loss\n",
+    "        imag_F_loss = self.alpha * (1 - imag_pt) ** self.gamma * imag_BCE_loss\n",
+    "\n",
+    "        if self.reduction == 'mean':\n",
+    "            return (torch.mean(real_F_loss) + torch.mean(imag_F_loss)) / 2\n",
+    "        elif self.reduction == 'sum':\n",
+    "            return torch.sum(real_F_loss) + torch.sum(imag_F_loss)\n",
+    "        else:\n",
+    "            return real_F_loss + imag_F_loss\n",
+    "\n",
+    "# Update the IoU calculation to handle complex values\n",
+    "def calculate_iou(pred, target, threshold=0.5):\n",
+    "    real_pred = (pred.real > threshold).float()\n",
+    "    imag_pred = (pred.imag > threshold).float()\n",
+    "    \n",
+    "    combined_pred = torch.logical_or(real_pred, imag_pred).float()\n",
+    "    \n",
+    "    intersection = (combined_pred * target).sum(dim=1)\n",
+    "    union = (combined_pred + target).sum(dim=1) - intersection\n",
+    "    iou = (intersection / union).mean().item()\n",
+    "    return iou"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "abb35ba2",
+   "metadata": {},
+   "source": [
+    "### Training with complex focal loss"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86d7526b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize and train the CResNet-18 model\n",
+    "model = ComplexResNet18().to(device)\n",
+    "criterion = ComplexFocalLoss()\n",
+    "\n",
+    "# Train the model and validate it\n",
+    "#0.001, 0.0001, 0.00001, 0.000001\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations =train_model(model, train_loader, valid_loader, criterion, initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3)\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd0c9d58",
+   "metadata": {},
+   "source": [
+    "### CVNN RV-BCE and CV-BCE Loss function implementation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99c736b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# RV BCE Loss Function Definition\n",
+    "class RealValuedBCELoss(nn.Module):\n",
+    "    def __init__(self, reduction='mean'):\n",
+    "        super(RealValuedBCELoss, self).__init__()\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        # Use only the real part of the complex inputs\n",
+    "        real_inputs = inputs.real\n",
+    "        BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction=self.reduction)\n",
+    "        return BCE_loss\n",
+    "\n",
+    "    \n",
+    "# CV BCE Loss Function Definition\n",
+    "class ComplexValuedBCELoss(nn.Module):\n",
+    "    def __init__(self, reduction='mean'):\n",
+    "        super(ComplexValuedBCELoss, self).__init__()\n",
+    "        self.reduction = reduction\n",
+    "\n",
+    "    def forward(self, inputs, targets):\n",
+    "        real_inputs = inputs.real\n",
+    "        imag_inputs = inputs.imag\n",
+    "\n",
+    "        # Calculate binary cross-entropy for both real and imaginary parts\n",
+    "        real_BCE_loss = F.binary_cross_entropy(real_inputs, targets, reduction=self.reduction)\n",
+    "        imag_BCE_loss = F.binary_cross_entropy(imag_inputs, targets, reduction=self.reduction)\n",
+    "        \n",
+    "        # Combine the losses (you can adjust the weighting if necessary)\n",
+    "        combined_BCE_loss = (real_BCE_loss + imag_BCE_loss) / 2\n",
+    "        return combined_BCE_loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d6930f39",
+   "metadata": {},
+   "source": [
+    "### RV-BCE Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9e59d4c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set the criterion for RV BCE\n",
+    "criterion = RealValuedBCELoss()\n",
+    "\n",
+    "# Train the ResNet-18 model with RV BCE\n",
+    "device = torch.device('cuda')\n",
+    "model = ComplexResNet18().to(device)\n",
+    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
+    "\n",
+    "# Start training with the previously defined train_model function\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations = train_model(\n",
+    "    model, train_loader, valid_loader, criterion, \n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001, 0.00001, 0.000001], num_epochs=50, patience=3\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93d19ea7",
+   "metadata": {},
+   "source": [
+    "### CV-BCE Training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c56d5b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set the criterion for CV BCE\n",
+    "criterion = ComplexValuedBCELoss()\n",
+    "\n",
+    "# Train the ResNet-18 model with CV BCE\n",
+    "device = torch.device('cuda')\n",
+    "model = ComplexResNet18().to(device)\n",
+    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
+    "\n",
+    "# Start training with the previously defined train_model function\n",
+    "model, train_losses, val_losses, val_accuracies, epoch_durations = train_model(\n",
+    "    model, train_loader, valid_loader, criterion, \n",
+    "    initial_lr=0.001, lr_steps=[0.001, 0.0001], num_epochs=50, patience=3\n",
+    ")\n",
+    "combined_epoch_time = sum(epoch_durations)\n",
+    "print(f\"Total time spent in epochs: {combined_epoch_time:.2f} seconds.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f4f6530e",
+   "metadata": {},
+   "source": [
+    "### Plot training result (Accuracy, loss vs epoch)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43676a01",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import json\n",
+    "import os\n",
+    "\n",
+    "# Ensure the directory exists\n",
+    "output_dir = 'cvnn_results/segmentation'\n",
+    "os.makedirs(output_dir, exist_ok=True)\n",
+    "\n",
+    "def save_metrics_to_json(train_losses, val_accuracies, epoch_durations, filename):\n",
+    "    \"\"\"\n",
+    "    Save the training losses and validation accuracies to a JSON file.\n",
+    "    \n",
+    "    Args:\n",
+    "        train_losses (list): List of training losses.\n",
+    "        val_accuracies (list): List of validation accuracies.\n",
+    "        filename (str): The file name for the JSON file.\n",
+    "    \"\"\"\n",
+    "    metrics = {\n",
+    "        \"train_losses\": train_losses,\n",
+    "        \"val_accuracies\": val_accuracies,\n",
+    "        \"epoch_durations\": epoch_durations\n",
+    "    }\n",
+    "    with open(os.path.join(output_dir, filename), 'w') as f:\n",
+    "        json.dump(metrics, f)\n",
+    "\n",
+    "def plot_training_metrics(train_losses, val_accuracies, plot_filename):\n",
+    "    \"\"\"\n",
+    "    Plot the training loss and validation accuracy, and mark the epoch where accuracy reaches 99%.\n",
+    "    \n",
+    "    Args:\n",
+    "        train_losses (list): List of training losses.\n",
+    "        val_accuracies (list): List of validation accuracies.\n",
+    "        plot_filename (str): The file name for saving the plot as SVG.\n",
+    "    \"\"\"\n",
+    "    epochs = range(1, len(train_losses) + 1)\n",
+    "\n",
+    "    plt.figure(figsize=(14, 6))\n",
+    "\n",
+    "    # Plot Training Loss\n",
+    "    plt.subplot(1, 2, 1)\n",
+    "    plt.plot(epochs, train_losses, label='Training Loss')\n",
+    "    plt.xlabel('Epochs')\n",
+    "    plt.ylabel('Loss')\n",
+    "    plt.title('Training Loss')\n",
+    "    plt.legend()\n",
+    "\n",
+    "    # Plot Validation Accuracy\n",
+    "    plt.subplot(1, 2, 2)\n",
+    "    plt.plot(epochs, val_accuracies, label='Validation Accuracy')\n",
+    "    plt.xlabel('Epochs')\n",
+    "    plt.ylabel('Accuracy (%)')\n",
+    "    plt.title('Validation Accuracy')\n",
+    "    plt.legend()\n",
+    "\n",
+    "    # Find the first epoch where validation accuracy reaches or exceeds 99%\n",
+    "    for i, acc in enumerate(val_accuracies):\n",
+    "        if acc >= 99:\n",
+    "            first_99_epoch = i + 1  # Epochs are 1-based\n",
+    "            plt.axvline(first_99_epoch, color='r', linestyle='--', label=f'99% reached at epoch {first_99_epoch}')\n",
+    "            break\n",
+    "\n",
+    "    plt.legend()\n",
+    "    plt.tight_layout()\n",
+    "\n",
+    "    # Save the plot as an SVG file\n",
+    "    plt.savefig(os.path.join(output_dir, plot_filename), format='svg')\n",
+    "    plt.show()\n",
+    "\n",
+    "# Save the metrics to JSON in cvnn_results/segmentation\n",
+    "save_metrics_to_json(train_losses, val_accuracies, epoch_durations, 'training_metrics.json')\n",
+    "\n",
+    "# Plot the metrics and highlight when accuracy reaches 99%, saving the plot as SVG\n",
+    "plot_training_metrics(train_losses, val_accuracies, 'training_metrics_plot.svg')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c6f4ea75",
+   "metadata": {},
+   "source": [
+    "### Evaluation "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a303080e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the pre-trained model for evaluation\n",
+    "import torch\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "\n",
+    "model_path = \"path/to/the/model\" #Please change this to the model path you trained\n",
+    "model = ComplexResNet18().to(device)\n",
+    "model.load_state_dict(torch.load(model_path, map_location=device))\n",
+    "model.eval()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0590b6ef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from tqdm import tqdm\n",
+    "from torch.utils.data import DataLoader\n",
+    "import numpy as np\n",
+    "\n",
+    "# Define thresholds for recall calculation\n",
+    "iou_thresholds = [0.5, 0.7, 0.9]\n",
+    "\n",
+    "# Initialize metrics\n",
+    "snr_results = {}\n",
+    "total_accuracy = 0.0\n",
+    "total_samples = 0\n",
+    "iou_scores = {th: 0.0 for th in iou_thresholds}\n",
+    "recall_counts = {th: 0 for th in iou_thresholds}\n",
+    "BATCH_SIZE = 64\n",
+    "# Create DataLoader for the entire dataset\n",
+    "full_dataset = WidebandSignalDataset(signal_ids=train + validation + test, return_snr=True)\n",
+    "full_loader = DataLoader(full_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6db6a18f",
+   "metadata": {},
+   "source": [
+    "### Bounding Box"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e396c72c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from collections import defaultdict\n",
+    "import time\n",
+    "from tqdm import tqdm\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "from scipy.optimize import linear_sum_assignment\n",
+    "\n",
+    "def expand_true(array, distance=1):\n",
+    "    # Create kernel of appropriate size\n",
+    "    kernel = torch.ones((1, 1, distance * 2 + 1), device=array.device)\n",
+    "    array = array.unsqueeze(1).float()  # Add channel dimension\n",
+    "    result = F.conv1d(array, kernel, padding=distance)\n",
+    "    result = result.squeeze(1)  # Remove the extra dimension\n",
+    "    \n",
+    "    # Convert values greater than 0 to `True`\n",
+    "    return result > 0\n",
+    "\n",
+    "# Define supporting functions based on your friend's code\n",
+    "def get_true_groups(tensor, device):\n",
+    "    assert tensor.dim() == 2, 'This function handles 2D tensor only'\n",
+    "    all_groups = []\n",
+    "    for i in range(tensor.size(0)):\n",
+    "        item = tensor[i]\n",
+    "        item = torch.cat([torch.tensor([False]).to(device), item, torch.tensor([False]).to(device)])\n",
+    "        diffs = item.float().diff()\n",
+    "        starts = (diffs == 1).nonzero(as_tuple=True)[0]\n",
+    "        ends = (diffs == -1).nonzero(as_tuple=True)[0] - 1\n",
+    "        groups = [(start.item(), end.item()) for start, end in zip(starts, ends)]\n",
+    "        all_groups.append(groups)\n",
+    "    return all_groups\n",
+    "\n",
+    "def get_target_boxes(metadata, number_of_bins, sample_rate=SAMPLE_RATE):\n",
+    "    scale_ratio = number_of_bins / sample_rate\n",
+    "    targets = []\n",
+    "    masks = torch.zeros(number_of_bins)\n",
+    "    for meta in metadata:\n",
+    "        f, b = meta['position']\n",
+    "        x1, x2 = math.floor((f-b/2)*scale_ratio), math.ceil((f+b/2)*scale_ratio)\n",
+    "        masks[x1:x2] = 1\n",
+    "        targets.append((x1, x2))\n",
+    "    return targets, masks\n",
+    "\n",
+    "def get_target_boxes_batch(batch_metadata, number_of_bins, sample_rate=SAMPLE_RATE):\n",
+    "    all_targets, all_masks = [], []\n",
+    "    for metadata in batch_metadata:\n",
+    "        targets, masks = get_target_boxes(metadata, number_of_bins, sample_rate)\n",
+    "        all_targets.append(targets)\n",
+    "        all_masks.append(masks)\n",
+    "    return all_targets, all_masks\n",
+    "\n",
+    "def calculate_iou(box1, box2):\n",
+    "    intersection = max(0, min(box1[1], box2[1]) - max(box1[0], box2[0]))\n",
+    "    union = max(box1[1], box2[1]) - min(box1[0], box2[0])\n",
+    "    return intersection / union if union != 0 else 0\n",
+    "\n",
+    "def match_targets(targets, preds):\n",
+    "    ious = []\n",
+    "    for target in targets:\n",
+    "        iou_targets = []\n",
+    "        for pred in preds:\n",
+    "            iou_targets.append(calculate_iou(target, pred))\n",
+    "        ious.append(iou_targets)\n",
+    "    return linear_sum_assignment(ious, maximize=True)\n",
+    "\n",
+    "def match_targets_batch(batch_targets, batch_preds):\n",
+    "    all_assignments = []\n",
+    "    for targets, preds in zip(batch_targets, batch_preds):\n",
+    "        all_assignments.append(match_targets(targets, preds))\n",
+    "    return all_assignments\n",
+    "\n",
+    "def calculate_matched_ious(target_boxes, prediction_boxes, matching):\n",
+    "    ious = [0 for _ in target_boxes]\n",
+    "    matching_dict = dict(zip(*matching))\n",
+    "    for target_index, target_box in enumerate(target_boxes):\n",
+    "        if target_index in matching_dict:\n",
+    "            box1 = target_box\n",
+    "            box2 = prediction_boxes[matching_dict[target_index]]\n",
+    "            ious[target_index] = calculate_iou(box1, box2)\n",
+    "    return ious\n",
+    "\n",
+    "def calculate_matched_iou_mean_batch(batch_target_boxes, batch_pred_boxes, batch_matching):\n",
+    "    all_ious = []\n",
+    "    for args in zip(batch_target_boxes, batch_pred_boxes, batch_matching):\n",
+    "        all_ious.append(calculate_matched_ious(*args))\n",
+    "    return all_ious\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24d483c1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import defaultdict\n",
+    "from tqdm import tqdm\n",
+    "def model_predictor(signals):\n",
+    "    # Use the already loaded model and apply thresholding\n",
+    "    signals = reshape_to_2d(signals)\n",
+    "    outputs = model(signals)\n",
+    "    return expand_true(outputs.real > 0.5)  # Use real part for thresholding\n",
+    "def evaluate(predictor, data_loader, device=\"cuda\"):\n",
+    "    snr_metrics = defaultdict(lambda: {\n",
+    "        \"iou_sum\": 0.0,\n",
+    "        \"iou_count\": 0,\n",
+    "        \"recall_counts\": defaultdict(int),\n",
+    "        \"total_samples\": defaultdict(int),\n",
+    "        \"correct_pixels\": 0,\n",
+    "        \"total_pixels\": 0\n",
+    "    })\n",
+    "    total_iou_sum, total_iou_count = 0.0, 0\n",
+    "    total_correct_pixels, total_total_pixels = 0, 0\n",
+    "    total_recall_counts = defaultdict(int)\n",
+    "    total_samples = defaultdict(int)\n",
+    "\n",
+    "    for inputs, masks, snrs_in_batch in tqdm(data_loader, desc=\"Evaluating\"):\n",
+    "        #inputs = inputs.to(device)\n",
+    "        inputs = reshape_to_2d(inputs).to(device)\n",
+    "        masks = masks.to(device)\n",
+    "        outputs = predictor(inputs)\n",
+    "\n",
+    "        for i in range(len(snrs_in_batch)):\n",
+    "            snr = snrs_in_batch[i].item()\n",
+    "            mask = masks[i]\n",
+    "            output = outputs[i]\n",
+    "\n",
+    "            # Ensure output matches mask shape\n",
+    "            if output.numel() != mask.numel():\n",
+    "                output = output.expand_as(mask) if output.numel() == 1 else output.reshape_as(mask)\n",
+    "\n",
+    "            thresholded_output = (output.real >= 0.5).float()\n",
+    "\n",
+    "            correct_pixels = (thresholded_output == mask).sum().item()\n",
+    "            total_pixels = mask.numel()\n",
+    "            snr_metrics[snr][\"correct_pixels\"] += correct_pixels\n",
+    "            snr_metrics[snr][\"total_pixels\"] += total_pixels\n",
+    "            total_correct_pixels += correct_pixels\n",
+    "            total_total_pixels += total_pixels\n",
+    "\n",
+    "            target_boxes = get_true_groups(mask.unsqueeze(0), device=device)[0]\n",
+    "            pred_boxes = get_true_groups(thresholded_output.unsqueeze(0), device=device)[0]\n",
+    "            if not target_boxes or not pred_boxes:\n",
+    "                continue\n",
+    "            matching = match_targets(target_boxes, pred_boxes)\n",
+    "            matched_ious = calculate_matched_ious(target_boxes, pred_boxes, matching)\n",
+    "\n",
+    "            snr_metrics[snr][\"iou_sum\"] += sum(matched_ious)\n",
+    "            snr_metrics[snr][\"iou_count\"] += len(matched_ious)\n",
+    "            total_iou_sum += sum(matched_ious)\n",
+    "            total_iou_count += len(matched_ious)\n",
+    "\n",
+    "            for th in iou_thresholds:\n",
+    "                true_positives = sum(1 for iou in matched_ious if iou >= th)\n",
+    "                snr_metrics[snr][\"recall_counts\"][th] += true_positives\n",
+    "                snr_metrics[snr][\"total_samples\"][th] += len(target_boxes)\n",
+    "                total_recall_counts[th] += true_positives\n",
+    "                total_samples[th] += len(target_boxes)\n",
+    "\n",
+    "    # Calculate overall metrics\n",
+    "    overall_accuracy = (total_correct_pixels / total_total_pixels) * 100 if total_total_pixels > 0 else 0\n",
+    "    overall_iou = total_iou_sum / total_iou_count if total_iou_count > 0 else 0\n",
+    "    overall_recall = {th: total_recall_counts[th] / total_samples[th] if total_samples[th] > 0 else 0 for th in iou_thresholds}\n",
+    "\n",
+    "    # Print overall results\n",
+    "    print(f\"Overall Accuracy: {overall_accuracy:.2f}%\")\n",
+    "    print(f\"Overall IoU Score: {overall_iou:.4f}\")\n",
+    "    for th in iou_thresholds:\n",
+    "        print(f\"Recall at threshold {th}: {overall_recall[th]:.4f}\")\n",
+    "\n",
+    "    # Print per-SNR results\n",
+    "    for snr, metrics in sorted(snr_metrics.items()):\n",
+    "        snr_accuracy = (metrics[\"correct_pixels\"] / metrics[\"total_pixels\"]) * 100 if metrics[\"total_pixels\"] > 0 else 0\n",
+    "        snr_iou = metrics[\"iou_sum\"] / metrics[\"iou_count\"] if metrics[\"iou_count\"] > 0 else 0\n",
+    "        print(f\"SNR: {snr} dB - Accuracy: {snr_accuracy:.2f}%\")\n",
+    "        print(f\"   IoU: {snr_iou:.4f}\")\n",
+    "        for th in iou_thresholds:\n",
+    "            recall = metrics[\"recall_counts\"][th] / metrics[\"total_samples\"][th] if metrics[\"total_samples\"][th] > 0 else 0\n",
+    "            print(f\"   Recall at threshold {th}: {recall:.4f}\")\n",
+    "\n",
+    "    return snr_metrics\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a71c18ba",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "snr_metrics = evaluate(model_predictor, full_loader, device=device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "87417c7b",
+   "metadata": {},
+   "source": [
+    "### Plot and Save"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1dbfb5e6",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import matplotlib.pyplot as plt\n",
+    "from pathlib import Path\n",
+    "\n",
+    "# Define the path for saving the JSON file and plots\n",
+    "save_path = Path(\"CMuSeNet_plots/Synthetic\")\n",
+    "save_path.mkdir(parents=True, exist_ok=True)\n",
+    "json_file_path = save_path / \"evaluation_results.json\"\n",
+    "\n",
+    "# Save metrics and plot results\n",
+    "def save_and_plot_results(snr_metrics, iou_thresholds):\n",
+    "    # Prepare data for plotting and JSON saving\n",
+    "    snr_values = sorted(snr_metrics.keys())\n",
+    "    iou_scores = [snr_metrics[snr][\"iou_sum\"] / snr_metrics[snr][\"iou_count\"] if snr_metrics[snr][\"iou_count\"] > 0 else 0 for snr in snr_values]\n",
+    "    accuracies = [(snr_metrics[snr][\"correct_pixels\"] / snr_metrics[snr][\"total_pixels\"]) * 100 if snr_metrics[snr][\"total_pixels\"] > 0 else 0 for snr in snr_values]\n",
+    "    recalls = {th: [(snr_metrics[snr][\"recall_counts\"][th] / snr_metrics[snr][\"total_samples\"][th]) if snr_metrics[snr][\"total_samples\"][th] > 0 else 0 for snr in snr_values] for th in iou_thresholds}\n",
+    "\n",
+    "    # Save results to JSON\n",
+    "    results = {\n",
+    "        \"SNR\": snr_values,\n",
+    "        \"IoU_Scores\": iou_scores,\n",
+    "        \"Accuracy\": accuracies,\n",
+    "        \"Recall\": {str(th): recalls[th] for th in iou_thresholds}\n",
+    "    }\n",
+    "    with open(json_file_path, \"w\") as f:\n",
+    "        json.dump(results, f, indent=4)\n",
+    "    print(f\"Results saved to {json_file_path}\")\n",
+    "\n",
+    "    # Plot IoU vs SNR\n",
+    "    plt.figure()\n",
+    "    plt.plot(snr_values, iou_scores, marker='o', label=\"IoU Score\")\n",
+    "    plt.xlabel(\"SNR (dB)\")\n",
+    "    plt.ylabel(\"IoU Score\")\n",
+    "    plt.title(\"IoU Score vs. SNR\")\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    plt.savefig(save_path / \"IoU_vs_SNR.png\")\n",
+    "    plt.savefig(save_path / \"IoU_vs_SNR.svg\")\n",
+    "    plt.show()\n",
+    "\n",
+    "    # Plot Accuracy vs SNR\n",
+    "    plt.figure()\n",
+    "    plt.plot(snr_values, accuracies, marker='o', label=\"Accuracy\")\n",
+    "    plt.xlabel(\"SNR (dB)\")\n",
+    "    plt.ylabel(\"Accuracy (%)\")\n",
+    "    plt.title(\"Accuracy vs. SNR (Threshold 0.5)\")\n",
+    "    plt.grid(True)\n",
+    "    plt.legend()\n",
+    "    plt.savefig(save_path / \"Accuracy_vs_SNR.png\")\n",
+    "    plt.savefig(save_path / \"Accuracy_vs_SNR.svg\")\n",
+    "    plt.show()\n",
+    "\n",
+    "    # Plot Recall vs SNR for each threshold\n",
+    "    for th in iou_thresholds:\n",
+    "        plt.figure()\n",
+    "        plt.plot(snr_values, recalls[th], marker='o', label=f\"Recall at {th}\")\n",
+    "        plt.xlabel(\"SNR (dB)\")\n",
+    "        plt.ylabel(\"Recall\")\n",
+    "        plt.title(f\"Recall vs. SNR (Threshold {th})\")\n",
+    "        plt.grid(True)\n",
+    "        plt.legend()\n",
+    "        plt.savefig(save_path / f\"Recall_vs_SNR_{th}.png\")\n",
+    "        plt.savefig(save_path / f\"Recall_vs_SNR_{th}.svg\")\n",
+    "        plt.show()\n",
+    "\n",
+    "# Call this after running evaluate() to save and plot results\n",
+    "save_and_plot_results(snr_metrics, iou_thresholds)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0c0d3e8",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

CMuSeNet_Synthetic_IQ_Generator/README.txt

@@ -0,0 +1,26 @@
+This matlab function is generated and tested in MATLAB 2022, and 2024
+
+Please open datagen.m script and run it with MATLAB to generate synthetic dataset with same configuration as CMuSeNet synthetic dataset.
+
+In this script you can change various setting such as channel (AWGN, Rician, Rayleigh), sample speed, range of SNR and sample bandwidth.
+
+This dataset is used to train CMuSeNet, complex-valued multi-signla segmentation Network.
+
+Please cite our paper if you use this dataset or synthetic dataset generation script.
+
+@inproceedings{shin2025cmusenet,
+  title={I Can't Believe It's Not Real: {CV-MuSeNet}: Complex-Valued Multi-Signal Segmentation},
+  author={Sangwon Shin and Mehmet C. Vuran},
+  booktitle={IEEE Dynamic Spectrum Access Networks (DySPAN)},
+  year={2025},
+  organization={IEEE}
+}
+
+Acknowledgement: Office of Naval Research, NSWC Crane N00174-23-1-0007
+This work relates to Department of Navy award N00174-23-1-0007 issued by the Office of Naval Research, NSWC Crane. Any opinions, 
+findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the Office of Naval Research.
+
+Following IQ samples generation script is coded by Prashant Subedi, Sangwon Shin and Dr. Mehmet Can Vuran - Cyber Physical Networking (CPN) Lab at University of Nebraska - Lincoln
+
+License:
+This IQ samples is licensed under the GPL family (General Public License) terms.

CMuSeNet_Synthetic_IQ_Generator/datagen.m

@@ -0,0 +1,28 @@
+path = "../diff-snr-matlab-simulated-data";
+
+for snr = -20:2:10
+    disp(snr);
+    mkdir(sprintf("%s/%d/", path, snr));
+    for i = 1:5000
+        name = string(i);
+        channelType = 'awgn'; %Supported channel type: awgn, rician (Flat), rayleigh (Flat)
+        [meta, data] = datagenWideband(snr, channelType);
+        split = reshape([real(data) imag(data)].', 1, []);
+        
+        % Save data file
+        mkdir(sprintf("%s/%d/%s", path, snr, name));
+        
+        
+        datafile = fopen(sprintf("%s/%d/%s/data.dat", path, snr, name), 'w');
+        fwrite(datafile, split, 'double');
+        fclose(datafile);
+        
+        % Save meta file
+        metafile = fopen(sprintf("%s/%d/%s/meta-data.json", path, snr, name), 'w');
+        fprintf(metafile, jsonencode(meta));
+        fclose(metafile);
+        
+        
+        disp(name);
+    end
+end

CMuSeNet_Synthetic_IQ_Generator/datagenTransmitter.m

@@ -0,0 +1,64 @@
+function transmittedSignal = datagenTransmitter( ...
+    modulation, ...
+    rolloffFactor, ...
+    filterSpanInSymbols, ...
+    samplesPerSymbol, ...
+    symbolRate, ...
+    messageDuration ...
+    )
+    requiresFilter = true;
+    if modulation == "QPSK"
+        bitsPerSymbol = 2;
+        modulator = comm.QPSKModulator( ...
+          'BitInput',                     true, ...
+          'PhaseOffset',                  pi/4, ...
+          'OutputDataType',               'double' ...
+        );
+    elseif modulation == "BPSK"
+        bitsPerSymbol = 1;
+        modulator = comm.BPSKModulator;
+    elseif modulation == "8-PSK"
+        bitsPerSymbol = 3;
+        modulator = @(x) qammod(bit2int(x, 3), 8);
+    elseif modulation == "8-QAM"
+        bitsPerSymbol = 3;
+        modulator = @(x) pskmod(bit2int(x, 3), 8);
+    elseif modulation == "16-QAM"
+        bitsPerSymbol = 4;
+        modulator = @(x) qammod(bit2int(x, 4), 16);
+    elseif modulation == "GMSK"
+        bitsPerSymbol = 1;
+        modulator = comm.GMSKModulator("SamplesPerSymbol", samplesPerSymbol, ...
+            "BitInput", true);
+        requiresFilter = false;
+   elseif modulation == "2-FSK"
+      bitsPerSymbol = 1;
+      fdev = floor(symbolRate/4);
+      samplesPerSymbol = 8;
+      modulator = @(x) fskmod(x, 2, fdev, samplesPerSymbol, symbolRate);
+      requiresFilter = false;
+    else
+        error("Not implemented " + modulation);
+    end    
+      
+     transmittedBin = randi( ...
+         [0 1], ...
+         bitsPerSymbol * symbolRate * messageDuration/samplesPerSymbol, ...
+         1 ...
+      );
+     
+     modulatedData = modulator(transmittedBin);        % Modulates the bits into QPSK symbols           
+
+    if requiresFilter
+        transmitterFilter = comm.RaisedCosineTransmitFilter( ...
+          'RolloffFactor',                rolloffFactor, ...
+          'FilterSpanInSymbols',          filterSpanInSymbols, ...
+          'OutputSamplesPerSymbol',       samplesPerSymbol ...
+        );
+        transmittedSignal = transmitterFilter(modulatedData); % Square root Raised Cosine Transmit Filter
+    else
+        transmittedSignal = modulatedData;
+    end
+    
+
+end

CMuSeNet_Synthetic_IQ_Generator/datagenWideband.m

@@ -0,0 +1,147 @@
+function [metadata, widebandSignal] = datagenWideband(SNRdB, fadingType)
+      % Constant for this function
+      RolloffFactor = 0.35;
+      RaisedCosineFilterSpan = 10;
+      Interpolation = 2;
+      NarrowBandBWs = [1e5, 2e5, 5e5, 1e6, 2e6];
+      WideBandBW = 20e6;
+      MaxSignals = 10;
+      %  Modulations = ["QPSK" "BPSK" "8-PSK" "8-QAM" "16-QAM" "2-FSK" ];
+      Modulations = ["QPSK" "BPSK" "8-PSK" "8-QAM" "16-QAM", "GMSK", "2-FSK"];
+      SamplingTime = 2/1000; % 2ms
+
+      TxPowerRange = [0, 20];
+      
+      numberOfSignals = randi([1, MaxSignals], 1);
+  
+      signalBW = randsample(NarrowBandBWs, numberOfSignals, true);
+      txPowers = randi(TxPowerRange, [numberOfSignals, 1]);
+
+      
+      minGap = 100e3; % 100kHz
+
+      maxBW = max(signalBW);
+
+      % Allocate a space for the frequencies
+      freqOffsets = [];
+      usedFreqs = [];
+      % A mechanism to prevent it from being stuck if there are too many
+      % wideband signals
+      maxLoops = numberOfSignals * 10;
+      % Generate non-overlapping frequencies
+      for i = 1:numberOfSignals
+        bw = signalBW(i);
+        % Generate a random frequency offset within the limits
+        while maxLoops > 0
+            maxLoops = maxLoops -  1; % prevent it from handing
+            freq = randi([-WideBandBW/2 + bw/2, WideBandBW/2 - bw/2]);
+            % Check if the frequency space for the new signal is already occupied or
+            % if the new signal is within minGap of an existing signal
+            overlap = false;
+            for j = 1:length(usedFreqs)
+                existing_bw = signalBW(j);
+                if abs(freq - usedFreqs(j)) < (bw + existing_bw)/2 + minGap
+                    overlap = true;
+                    break;
+                end
+            end
+            if ~overlap
+                % If not, add the frequency to the used frequencies and break the loop
+                usedFreqs = [usedFreqs freq];
+                freqOffsets = [freqOffsets freq];
+                break
+            end
+            % If the frequency space is occupied or too close to another signal,
+            % generate a new random frequency
+        end
+
+        if maxLoops <= 0
+            numberOfSignals = length(freqOffsets);
+            disp("Stopping because couldn't place signal");
+            disp(signalBW);
+            break;
+        end
+
+      end
+
+
+      signals = [];
+      metadata = [];
+
+      lowestPowerSignal = min(txPowers);
+      noisePower = min(txPowers) - SNRdB;
+
+      
+      for i = 1: numberOfSignals
+          modulation = randsample(Modulations, 1);
+          txPower = txPowers(i);
+          bw = signalBW(i);
+          % Should the divisor be 20 ?
+          signal = datagenTransmitter( ...
+              modulation, ...
+              RolloffFactor, ...
+              RaisedCosineFilterSpan, ...
+              Interpolation, ...
+              bw, ...
+              SamplingTime...811
+           );
+          
+          % Scale the signal
+          signal = signal/sqrt(mean(abs(signal).^2));
+
+          % Scale to correct power
+          signal = 10^(txPower/20)*signal;
+
+          pwr = 10*log10(mean(abs(signal).^2));
+          
+          
+          if bw ~= maxBW
+              signal = resample(signal, maxBW/1e5, bw/1e5);
+          end
+          signals = [signals signal];
+          metadata = [metadata; struct("fc", freqOffsets(i), "bw", bw, "mod", modulation, "txPower", txPower, "noisePower", noisePower)];
+
+      end
+      mbc = comm.MultibandCombiner( ...
+        InputSampleRate=maxBW, ...
+        FrequencyOffsets=freqOffsets, ...
+        OutputSampleRateSource="property", ...
+        OutputSampleRate=WideBandBW ...
+      );
+      
+      combinedsig = mbc(signals);
+      % Channel configuration
+      fd = 30; % Max Doppler shift in Hz
+      Ts = 1/WideBandBW; % Sampling time
+      chan = [];
+        
+      switch lower(fadingType)
+          case 'awgn'
+              % Just noise without fading
+              widebandSignal = awgn(combinedsig, SNRdB, lowestPowerSignal);
+    
+          case 'rayleigh'
+              rayleighChan = comm.RayleighChannel( ...
+                  'SampleRate', WideBandBW, ...
+                  'PathDelays', 0, ...
+                  'AveragePathGains', 0, ...
+                  'MaximumDopplerShift', 30 ...
+              );
+              fadedSignal = rayleighChan(combinedsig);
+              widebandSignal = awgn(fadedSignal, SNRdB, lowestPowerSignal);  % Add AWGN
+    
+          case 'rician'
+              ricianChan = comm.RicianChannel( ...
+                  'SampleRate', WideBandBW, ...
+                  'PathDelays', 0, ...
+                  'AveragePathGains', 0, ...
+                  'KFactor', 10, ...
+                  'MaximumDopplerShift', 30 ...
+              );
+              fadedSignal = ricianChan(combinedsig);
+              widebandSignal = awgn(fadedSignal, SNRdB, lowestPowerSignal);  % Add AWGN
+    
+          otherwise
+              error('Unsupported fading type: %s', fadingType);
+      end
+end