{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5a13ad6b-56c9-4381-b376-1765f6dd7553",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Импортирование библиотек"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7311cb4a-5bf3-4268-b431-43eea10e9ed6",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "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": "markdown",
   "id": "384de097-82c6-41f5-bda9-b2f54bc99593",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Подготовка и обучение детектирование"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "46e4dc99-6994-4fee-a32e-f3983bd991bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_and_learning_detection(num_classes, num_samples, path_dataset, model_name, config_name, model):\n",
    "    num_samples_per_class = num_samples // num_classes\n",
    "\n",
    "    #----------Создаём папку для сохранения результатов обучения--------------\n",
    "    \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",
    "    for subdir in subdirs:\n",
    "        files = os.listdir(path_dataset + subdir + '/')\n",
    "        num_samples_per_class = min(num_samples_per_class, len(files))\n",
    "    for subdir in subdirs:\n",
    "        files = os.listdir(path_dataset + subdir + '/')\n",
    "        random.shuffle(files)\n",
    "        files_to_process = files[:num_samples_per_class]\n",
    "        for file in files_to_process:\n",
    "            row = pd.DataFrame({pd_columns[0]: [str(path_dataset + subdir + '/' + file)]})\n",
    "            df = pd.concat([df, row], ignore_index=True)\n",
    "    \n",
    "    df.to_csv(path_res + 'dataset.csv', index=False)\n",
    "    \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(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",
    "            image_real = np.asarray(cv2.split(cv2.imread(self.sig_filenames[idx][:-8]+'real.jpg')), dtype=np.float32)\n",
    "            image_imag = np.asarray(cv2.split(cv2.imread(self.sig_filenames[idx][:-8]+'imag.jpg')), dtype=np.float32)\n",
    "            image_spec = np.asarray(cv2.split(cv2.imread(self.sig_filenames[idx][:-8]+'spec.jpg')), dtype=np.float32)\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, image_spec, 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, img3, label) in tqdm(dataloaders[phase]):\n",
    "                img1, img2, img3, label = img1.to(device), img2.to(device), img3.to(device), label.to(device)\n",
    "                optimizer.zero_grad()\n",
    "    \n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    output = model([img1, img2, img3])\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",
    "    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": "markdown",
   "id": "93c136ee",
   "metadata": {},
   "source": [
    "### Ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52e8d4c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\snytk\\miniconda3\\envs\\python311\\Lib\\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",
      "C:\\Users\\snytk\\miniconda3\\envs\\python311\\Lib\\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": [
      "Epoch : 1\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████| 658/658 [1:00:26<00:00,  5.51s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.6663 Acc: 0.9241\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 282/282 [02:45<00:00,  1.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.4023 Acc: 0.9557\n",
      "\n",
      "Epoch : 2\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 658/658 [43:11<00:00,  3.94s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Loss: 0.4096 Acc: 0.9514\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 282/282 [00:47<00:00,  5.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val Loss: 0.3390 Acc: 0.9574\n",
      "\n",
      "Epoch : 3\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|██████▍                                                                          | 52/658 [04:26<51:43,  5.12s/it]"
     ]
    }
   ],
   "source": [
    "#----------Инициализируем модель и параметры обучения--------------\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "cv2.destroyAllWindows()\n",
    "gc.collect()\n",
    "\n",
    "num_classes = 3\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",
    "model3 = models.resnet101(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",
    "model3.fc = nn.Linear(model3.fc.in_features, num_classes)\n",
    "\n",
    "class Ensemble(nn.Module):\n",
    "    def __init__(self, model1, model2, model3):\n",
    "        super(Ensemble, self).__init__()\n",
    "        self.model1 = model1\n",
    "        self.model2 = model2\n",
    "        self.model3 = model3\n",
    "        self.fc = nn.Linear(3 * num_classes, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.model1(x[0])\n",
    "        x2 = self.model2(x[1])\n",
    "        x3 = self.model3(x[2])\n",
    "        x = torch.cat((x1, x2, x3), dim=1)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "\n",
    "model = Ensemble(model1, model2, model3)\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 = 5000, path_dataset = \"C:/Users/snytk/Lerning_NN_for_work/datasets_jpg/915_jpg_learning/\", \n",
    "                     model_name = config_name+\"_915_jpg_\", config_name = config_name, model=model)\n",
    "\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "cv2.destroyAllWindows()\n",
    "del model\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d18676",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eab69324",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Отсутствует",
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}