notebooks/Datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eee8008a-5745-48e5-8d8a-d219514757aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_directory = \"data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/\"\n",
    "\n",
    "# I/Q dataset paramters\n",
    "num_iq_samples = 1024\n",
    "only_use_start_of_burst = True\n",
    "\n",
    "# Spectrogram parameters\n",
    "n_fft = 1024\n",
    "time_dim = 512\n",
    "n_samples = n_fft * time_dim\n",
    "yolo_class_list = [\"wifi\", \"anom_wifi\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da50772c-a1fb-4946-ade0-5f63a1bde799",
   "metadata": {},
   "source": [
    "# I/Q dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "12f25ae3-67ad-41e7-9a2d-da3c206518b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class List: ['wifi', 'anom_wifi']\n",
      "\n",
      "Dataset info:\n",
      "len(dataset)=1807, class_counts={'wifi': 903, 'anom_wifi': 904}\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torchsig.transforms as ST\n",
    "from sigmf_pytorch_dataset import SigMFDataset\n",
    "\n",
    "transform = ST.Compose([\n",
    "    ST.Normalize(norm=np.inf),\n",
    "    ST.ComplexTo2D(),\n",
    "])\n",
    "\n",
    "dataset = SigMFDataset( \n",
    "    root=data_directory, \n",
    "    sample_count= num_iq_samples, \n",
    "    transform = transform,\n",
    "    class_list = yolo_class_list,\n",
    "    only_first_samples=only_use_start_of_burst,\n",
    ")\n",
    "\n",
    "class_counts = {dataset.class_list[k]:v for k,v in dataset.get_class_counts().items()}\n",
    "print(\"\\nDataset info:\")\n",
    "print(f\"{len(dataset)=}, {class_counts=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5620b777-3fe7-4ba1-afb7-64962438e2f9",
   "metadata": {},
   "source": [
    "# Spectrogram dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4ab3fe50-0a3f-450e-b1b9-3e514daf8732",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import yaml\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "\n",
    "from data import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c48d5fe-5b19-45de-95b0-8ef5d42f84ee",
   "metadata": {},
   "source": [
    "# Convert SigMF to spectrograms and YOLO labels (only needs to run once)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "84f435b9-9228-4d94-a88d-3cc95391984f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                    | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loaded \n",
      " Data file: data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/anomwifi-encoding-0-sendsps-10000000.s16.zst.sigmf-data \n",
      " SigMF-Meta file: data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/anomwifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n",
      "Saving data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/anomwifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n",
      "Writing YOLO txt label files\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|                                                                                 | 0/1409 [00:00<?, ?it/s]\u001b[A\n",
      "  7%|████▉                                                                  | 98/1409 [00:00<00:01, 972.02it/s]\u001b[A\n",
      " 14%|█████████▉                                                            | 199/1409 [00:00<00:01, 994.23it/s]\u001b[A\n",
      " 23%|███████████████▌                                                     | 319/1409 [00:00<00:01, 1084.34it/s]\u001b[A\n",
      " 30%|████████████████████▉                                                | 428/1409 [00:00<00:00, 1046.69it/s]\u001b[A\n",
      " 38%|██████████████████████████▍                                           | 533/1409 [00:00<00:00, 960.76it/s]\u001b[A\n",
      " 45%|███████████████████████████████▎                                      | 631/1409 [00:00<00:00, 959.59it/s]\u001b[A\n",
      "100%|████████████████████████████████████████████████████████████████████| 1409/1409 [00:00<00:00, 1699.71it/s]\u001b[A\n",
      " 50%|██████████████████████████████████████                                      | 1/2 [00:04<00:04,  4.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/anomwifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n",
      "\n",
      "Loaded \n",
      " Data file: data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/wifi-encoding-0-sendsps-10000000.s16.zst.sigmf-data \n",
      " SigMF-Meta file: data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/wifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n",
      "Saving data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/wifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n",
      "Writing YOLO txt label files\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|                                                                                 | 0/1408 [00:00<?, ?it/s]\u001b[A\n",
      "  6%|████▍                                                                  | 87/1408 [00:00<00:01, 864.53it/s]\u001b[A\n",
      " 15%|██████████▍                                                          | 213/1408 [00:00<00:01, 1090.79it/s]\u001b[A\n",
      " 24%|████████████████▋                                                    | 340/1408 [00:00<00:00, 1168.18it/s]\u001b[A\n",
      " 32%|██████████████████████▍                                              | 457/1408 [00:00<00:00, 1143.65it/s]\u001b[A\n",
      " 41%|████████████████████████████                                         | 572/1408 [00:00<00:00, 1086.58it/s]\u001b[A\n",
      " 49%|█████████████████████████████████▌                                   | 684/1408 [00:00<00:00, 1095.28it/s]\u001b[A\n",
      "100%|████████████████████████████████████████████████████████████████████| 1408/1408 [00:00<00:00, 1832.87it/s]\u001b[A\n",
      "100%|████████████████████████████████████████████████████████████████████████████| 2/2 [00:08<00:00,  4.27s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving data/gamutrf/gamutrf-sd-gr-ieee-wifi/test_offline/wifi-encoding-0-sendsps-10000000.s16.zst.sigmf-meta\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "labels = set()\n",
    "image_dirs = []\n",
    "\n",
    "for f in tqdm(glob.glob(str(Path(data_directory,\"*-meta\")))):\n",
    "    \n",
    "    d = Data(f)\n",
    "    d.generate_spectrograms(n_samples, n_fft, cmap_str=\"turbo\", overwrite=False)\n",
    "    yolo_label_outdir = str(Path(data_directory,\"yolo\",\"labels\"))\n",
    "    yolo_image_outdir = str(Path(data_directory,\"yolo\",\"images\"))\n",
    "    image_dirs.append(yolo_image_outdir)\n",
    "    d.export_yolo(yolo_label_outdir, image_outdir=yolo_image_outdir, yolo_class_list=yolo_class_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a246b77-1b5e-49d2-a6d7-fb1def0f9635",
   "metadata": {},
   "source": [
    "# Creat YOLOv8 data.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4835ab38-fcee-4122-80d6-2665f3840b65",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_dirs = list(set(image_dirs))\n",
    "\n",
    "\n",
    "yolo_yaml = {}\n",
    "yolo_yaml[\"train\"] = image_dirs[0]\n",
    "yolo_yaml[\"names\"] = yolo_class_list\n",
    "yolo_yaml[\"nc\"] = len(yolo_class_list)\n",
    "\n",
    "with open('data.yml', 'w') as outfile:\n",
    "    yaml.dump(yolo_yaml, outfile, default_flow_style=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8b4ecba3-89bf-4aa5-8a5b-8ff80e50e1aa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ebc093-fa3c-472c-95bd-63365d2bcb9d",
   "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
}