{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e1db882b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sibscience-4/from_ssh/DroneDetector/.venv-train/lib/python3.12/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.\n",
      "  warnings.warn(\"Unable to import Axes3D. This may be due to multiple versions of \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "191"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import default_generator, randperm\n",
    "from torch.utils.data.dataset import Subset\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.io import read_image\n",
    "from importlib import import_module\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import models\n",
    "import torch, torchvision\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import os, shutil\n",
    "import mlconfig\n",
    "import random\n",
    "import shutil\n",
    "import timeit\n",
    "import copy\n",
    "import time\n",
    "import cv2\n",
    "import csv\n",
    "import sys\n",
    "import io\n",
    "import gc\n",
    "\n",
    "plt.rcParams[\"savefig.bbox\"] = 'tight'\n",
    "torch.manual_seed(1)\n",
    "#matplotlib.use('Agg')\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(device)\n",
    "torch.cuda.empty_cache()\n",
    "cv2.destroyAllWindows()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8e009995",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_and_learning_detection(num_classes, num_samples, path_dataset, model_name, config_name, model,selected_freq):\n",
    "    num_samples_per_class = num_samples // num_classes\n",
    "\n",
    "    #----------Создаём папку для сохранения результатов обучения--------------\n",
    "    os.makedirs(\"models\", exist_ok=True)\n",
    "    ind = 1\n",
    "    while True:\n",
    "        if os.path.exists(\"models/\" + model_name + str(ind)):\n",
    "            ind += 1\n",
    "        else:\n",
    "            os.mkdir(\"models/\" + model_name + str(ind))\n",
    "            path_res = \"models/\" + model_name + str(ind) + '/'\n",
    "            break\n",
    "    \n",
    "    #----------Создаём файл dataset.csv для обучения--------------\n",
    "    \n",
    "    pd_columns = ['file_name']\n",
    "    df = pd.DataFrame(columns=pd_columns)\n",
    "    \n",
    "    subdirs = os.listdir(path_dataset)\n",
    "    \n",
    "    for subdir in subdirs:\n",
    "        freq_dir = os.path.join(path_dataset, subdir, str(selected_freq)+\"_jpg\")\n",
    "        if not os.path.isdir(freq_dir):\n",
    "            print(\"Error1\")\n",
    "            continue\n",
    "        \n",
    "        files_k=[f for f in os.listdir(freq_dir)]\n",
    "        print(len(files_k))\n",
    "        \n",
    "        files = [\n",
    "                f for f in os.listdir(freq_dir)\n",
    "                if os.path.isfile(os.path.join(freq_dir, f)) and f.endswith('imag.png')\n",
    "            ]\n",
    "        num_samples_per_class = min(num_samples_per_class, len(files))\n",
    "        print(f\"num_samples per class {subdir} is {num_samples_per_class}\")\n",
    "\n",
    "    for subdir in subdirs:\n",
    "        freq_dir = os.path.join(path_dataset, subdir, str(selected_freq)+\"_jpg\")\n",
    "        if not os.path.isdir(freq_dir):\n",
    "            print(\"Error1\")\n",
    "            continue\n",
    "\n",
    "        files = [\n",
    "            f for f in os.listdir(freq_dir)\n",
    "            if os.path.isfile(os.path.join(freq_dir, f)) and f.endswith('imag.png')\n",
    "        ]\n",
    "        random.shuffle(files)\n",
    "        files_to_process = files[:num_samples_per_class]\n",
    "\n",
    "        for file in files_to_process:\n",
    "            row = pd.DataFrame({\n",
    "                pd_columns[0]: [str(os.path.join(freq_dir, file))]\n",
    "            })\n",
    "            df = pd.concat([df, row], ignore_index=True)\n",
    "\n",
    "    dataset_csv_path = os.path.join(path_res, 'dataset.csv')\n",
    "    df.to_csv(dataset_csv_path, index=False)\n",
    "\n",
    "    if not os.path.exists(dataset_csv_path):\n",
    "        raise RuntimeError(f'dataset.csv was not created: {dataset_csv_path}')\n",
    "    #----------Импортируем параметры для обучения--------------\n",
    "    \n",
    "    def load_function(attr):\n",
    "        module_, func = attr.rsplit('.', maxsplit=1)\n",
    "        return getattr(import_module(module_), func)\n",
    "    \n",
    "    config = mlconfig.load('config_' + config_name + '.yaml')\n",
    "    \n",
    "    #----------Создаём класс датасета--------------\n",
    "    \n",
    "    class MyDataset(Dataset):\n",
    "        def __init__(self, path_dataset, csv_file):\n",
    "            data=[]\n",
    "            with open(os.path.join(path_dataset, csv_file), newline='') as csvfile:\n",
    "                reader = csv.reader(csvfile, delimiter=' ', quotechar='|')\n",
    "                for row in list(reader)[1:]:\n",
    "                    row = str(row)\n",
    "                    data.append(row[2: len(row)-2])\n",
    "            self.sig_filenames = data\n",
    "            self.path_dataset = path_dataset\n",
    "    \n",
    "        def __len__(self):\n",
    "            return len(self.sig_filenames)\n",
    "    \n",
    "        def __getitem__(self, idx):\n",
    "            base = os.path.splitext(self.sig_filenames[idx])[0]\n",
    "            if base[-4:]==\"real\":\n",
    "                image_real = np.asarray(cv2.split(cv2.imread(base + '.png')), dtype=np.float32)\n",
    "                image_imag = np.asarray(cv2.split(cv2.imread(base[:-4] + 'imag.png')), dtype=np.float32)\n",
    "            if base[-4:]==\"imag\":\n",
    "                image_real = np.asarray(cv2.split(cv2.imread(base[:-4] + 'real.png')), dtype=np.float32)\n",
    "                image_imag = np.asarray(cv2.split(cv2.imread(base + '.png')), dtype=np.float32)\n",
    "                \n",
    "                \n",
    "            \n",
    "            if 'drone' in list(self.sig_filenames[idx].split('/')):\n",
    "                label = torch.tensor(0)\n",
    "            if 'noise' in list(self.sig_filenames[idx].split('/')):\n",
    "                label = torch.tensor(1)\n",
    "            return image_real, image_imag, label\n",
    "    \n",
    "    #----------Создаём датасет--------------\n",
    "    \n",
    "    dataset = MyDataset(path_dataset=path_res, csv_file='dataset.csv')\n",
    "    train_set, valid_set = torch.utils.data.random_split(dataset, [0.7, 0.3], generator=torch.Generator().manual_seed(42))\n",
    "    batch_size = config.batch_size\n",
    "    train_dataloader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "    valid_dataloader = torch.utils.data.DataLoader(valid_set, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "    \n",
    "    dataloaders = {}\n",
    "    dataloaders['train'] = train_dataloader\n",
    "    dataloaders['val'] = valid_dataloader\n",
    "    dataset_sizes = {}\n",
    "    dataset_sizes['train'] = len(train_set)\n",
    "    dataset_sizes['val'] = len(valid_set)\n",
    "\n",
    "    #----------Обучаем модель--------------\n",
    "\n",
    "    val_loss = []\n",
    "    val_acc = []\n",
    "    train_loss = []\n",
    "    train_acc = []\n",
    "    epochs = config.epoch\n",
    "    \n",
    "    best_acc = 0.0\n",
    "    best_model = copy.deepcopy(model.state_dict())\n",
    "    limit = config.limit\n",
    "    epoch_limit = epochs\n",
    "    \n",
    "    start = timeit.default_timer()\n",
    "    for epoch in range(1, epochs+1):\n",
    "        print(f\"Epoch : {epoch}\\n\")\n",
    "        dataloader = None\n",
    "    \n",
    "        for phase in ['train', 'val']:\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "    \n",
    "            for (img1, img2, label) in tqdm(dataloaders[phase]):\n",
    "                img1, img2, label = img1.to(device), img2.to(device), label.to(device)\n",
    "                optimizer.zero_grad()\n",
    "    \n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    output = model([img1, img2])\n",
    "                    _, pred = torch.max(output.data, 1)\n",
    "                    loss = criterion(output, label)\n",
    "                    if phase=='train' :\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "                \n",
    "                running_loss += loss.item() * 3 * img1.size(0)\n",
    "                running_corrects += torch.sum(pred == label.data)\n",
    "    \n",
    "            epoch_loss = running_loss / dataset_sizes[phase]\n",
    "            epoch_acc = running_corrects.double() / dataset_sizes[phase]\n",
    "                \n",
    "            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))\n",
    "                    \n",
    "            if phase=='train' :\n",
    "                train_loss.append(epoch_loss)\n",
    "                train_acc.append(epoch_acc)\n",
    "            else :\n",
    "                val_loss.append(epoch_loss)\n",
    "                val_acc.append(epoch_acc)\n",
    "                if val_acc[-1] > best_acc :\n",
    "                    ind_limit = 0\n",
    "                    best_acc = val_acc[-1]\n",
    "                    best_model = copy.deepcopy(model.state_dict())\n",
    "                    torch.save(best_model, path_res + model_name + '.pth')\n",
    "                else:\n",
    "                    ind_limit += 1\n",
    "                    \n",
    "                if ind_limit >= limit:\n",
    "                    break\n",
    "                \n",
    "        if ind_limit >= limit:\n",
    "            epoch_limit = epoch\n",
    "            break\n",
    "            \n",
    "        print()\n",
    "    \n",
    "    end = timeit.default_timer()\n",
    "    print(f\"Total time elapsed = {end - start} seconds\")\n",
    "    epoch_limit += 1\n",
    "    \n",
    "    #----------Вывод графиков и сохранение результатов обучения--------------\n",
    "    \n",
    "    train_acc = np.asarray(list(map(lambda x: x.item(), train_acc)))\n",
    "    val_acc = np.asarray(list(map(lambda x: x.item(), val_acc)))\n",
    "    \n",
    "    np.save(path_res+'train_acc.npy', train_acc)\n",
    "    np.save(path_res+'val_acc.npy', val_acc)\n",
    "    np.save(path_res+'train_loss.npy', train_loss)\n",
    "    np.save(path_res+'val_loss.npy', val_loss)\n",
    "    \n",
    "    plt.figure()\n",
    "    plt.plot(range(1,epoch_limit), train_loss, color='blue')\n",
    "    plt.plot(range(1,epoch_limit), val_loss, color='red')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')  \n",
    "             \n",
    "    plt.title('Loss Curve')\n",
    "    plt.legend(['Train Loss', 'Validation Loss'])\n",
    "    plt.show()\n",
    "    plt.clf()\n",
    "    plt.cla()\n",
    "    plt.close()\n",
    "    \n",
    "    plt.figure()\n",
    "    plt.plot(range(1,epoch_limit), train_acc, color='blue')\n",
    "    plt.plot(range(1,epoch_limit), val_acc, color='red')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('Accuracy Curve')\n",
    "    plt.legend(['Train Accuracy', 'Validation Accuracy'])\n",
    "    plt.show()\n",
    "    \n",
    "    plt.clf()\n",
    "    plt.cla()\n",
    "    plt.close()\n",
    "    torch.cuda.empty_cache()\n",
    "    cv2.destroyAllWindows()\n",
    "    del model\n",
    "    gc.collect()\n",
    "\n",
    "    return path_res, model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53717d63",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sibscience-4/from_ssh/DroneDetector/.venv-train/lib/python3.12/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/sibscience-4/from_ssh/DroneDetector/.venv-train/lib/python3.12/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9520\n",
      "num_samples per class noise is 2380\n",
      "5060\n",
      "num_samples per class drone is 1265\n",
      "Epoch : 1\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 442/442 [01:44<00:00,  4.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.1155 Acc: 0.9870\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 189/189 [00:21<00:00,  8.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0067 Acc: 0.9960\n",
      "\n",
      "Epoch : 2\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 442/442 [01:45<00:00,  4.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0032 Acc: 0.9983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 189/189 [00:22<00:00,  8.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0027 Acc: 0.9960\n",
      "\n",
      "Epoch : 3\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 442/442 [01:45<00:00,  4.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0014 Acc: 0.9983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 189/189 [00:22<00:00,  8.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0007 Acc: 0.9960\n",
      "\n",
      "Epoch : 4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 442/442 [01:46<00:00,  4.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0009 Acc: 0.9983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 189/189 [00:22<00:00,  8.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0007 Acc: 0.9960\n",
      "\n",
      "Epoch : 5\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████| 442/442 [01:46<00:00,  4.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0005 Acc: 0.9983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 189/189 [00:21<00:00,  8.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0005 Acc: 0.9960\n",
      "\n",
      "Epoch : 6\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 442/442 [01:46<00:00,  4.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.0004 Acc: 0.9983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 189/189 [00:22<00:00,  8.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.0003 Acc: 0.9960\n",
      "Total time elapsed = 769.277186642983 seconds\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27540 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 31344 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "67"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#----------Инициализируем модель и параметры обучения--------------\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "cv2.destroyAllWindows()\n",
    "gc.collect()\n",
    "\n",
    "config_name = \"ensemble\"\n",
    "    \n",
    "def load_function(attr):\n",
    "    module_, func = attr.rsplit('.', maxsplit=1)\n",
    "    return getattr(import_module(module_), func)\n",
    "    \n",
    "config = mlconfig.load('config_' + config_name + '.yaml')\n",
    "\n",
    "model1 = models.resnet18(pretrained=False)\n",
    "model2 = models.resnet50(pretrained=False)\n",
    "\n",
    "num_classes = 2\n",
    "\n",
    "model1.fc = nn.Linear(model1.fc.in_features, num_classes)\n",
    "model2.fc = nn.Linear(model2.fc.in_features, num_classes)\n",
    "\n",
    "class Ensemble(nn.Module):\n",
    "    def __init__(self, model1, model2):\n",
    "        super(Ensemble, self).__init__()\n",
    "        self.model1 = model1\n",
    "        self.model2 = model2\n",
    "        self.fc = nn.Linear(2 * num_classes, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.model1(x[0])\n",
    "        x2 = self.model2(x[1])\n",
    "        x = torch.cat((x1, x2), dim=1)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "model = Ensemble(model1, model2)\n",
    "\n",
    "optimizer = load_function(config.optimizer.name)(model.parameters(), lr=config.optimizer.lr)\n",
    "criterion = load_function(config.loss_function.name)()\n",
    "scheduler = load_function(config.scheduler.name)(optimizer, step_size=config.scheduler.step_size, gamma=config.scheduler.gamma)\n",
    "\n",
    "if device != 'cpu':\n",
    "    model = model.to(device)\n",
    "\n",
    "#----------Создания датасета и обучение модели--------------\n",
    "\n",
    "path_res, model_name = prepare_and_learning_detection(num_classes = num_classes, num_samples = 10000, path_dataset = \"/mnt/data/Dataset_img\", \n",
    "                     selected_freq=2400,model_name = config_name+\"2400_\", config_name = config_name, model=model)\n",
    "\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "cv2.destroyAllWindows()\n",
    "del model\n",
    "gc.collect()"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
