{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4fdb98fc-65bb-467e-be0c-168fee9b0fca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<contextlib.ExitStack at 0x20af0408250>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import io\n",
    "import cv2\n",
    "import copy\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import Dataset\n",
    "from torch import default_generator, randperm\n",
    "from PIL import Image\n",
    "#from torch._utils import _accumulate\n",
    "import csv\n",
    "from torch.utils.data.dataset import Subset\n",
    "from scipy import ndimage\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "batch_size = 16\n",
    "momentum=0.9\n",
    "lr = 1e-3\n",
    "import random\n",
    "sub_sample = 0.5\n",
    "import matplotlib\n",
    "import gc\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['agg.path.chunksize'] = 256*256\n",
    "plt.ioff()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9267fbe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sig2pic_with_spec(path_to_data, filename, figsize=(16,16), dpi=16, resize = None):\n",
    "    try:\n",
    "        array = np.load(path_to_data+filename)\n",
    "\n",
    "        if resize != None:\n",
    "            resized_real = cv2.resize(array[0], resize)\n",
    "            resized_mag = cv2.resize(array[1], resize)\n",
    "            resized_spectr = cv2.resize(array[2], resize)\n",
    "            img = np.asarray(array, dtype=np.float32)\n",
    "            return img\n",
    "        \n",
    "        img = np.asarray(array, dtype=np.float32)\n",
    "        return img\n",
    "    \n",
    "    except Exception as e:\n",
    "        print(str(e))\n",
    "        return None\n",
    "    \n",
    "    \n",
    "    \n",
    "\n",
    "def pre_func_ensemble(data=None, src ='', ind_inference=0):\n",
    "    try:\n",
    "        import matplotlib.pyplot as plt\n",
    "        matplotlib.use('Agg')\n",
    "        plt.ioff()\n",
    "\n",
    "        figsize = (16, 16)\n",
    "        dpi = 16\n",
    "\n",
    "        signal = np.vectorize(complex)(data[0], data[1])\n",
    "        #np.save(src + '_inference_2400_' + str(ind_inference) + '.npy', signal)\n",
    "        spec = transform.Spectrogram(nperseg=256)\n",
    "        spectr = np.array(spec(signal)[:,:figsize[0] * dpi])\n",
    "        fig1 = plt.figure(figsize = figsize)\n",
    "        plt.axes(ylim=(-1, 1))\n",
    "        sigr = signal.real\n",
    "        sigi = signal.imag\n",
    "        \n",
    "        plt.plot(sigr, color='black')\n",
    "        plt.gca().set_axis_off()\n",
    "        plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)\n",
    "        plt.margins(0,0)\n",
    "        buf1 = io.BytesIO()\n",
    "        fig1.savefig(buf1, format=\"png\", dpi=dpi)\n",
    "        buf1.seek(0)\n",
    "        img_arr1 = np.frombuffer(buf1.getvalue(), dtype=np.uint8)\n",
    "        buf1.close()\n",
    "        img1 = cv2.imdecode(img_arr1, 1)\n",
    "        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)\n",
    "        plt.clf()\n",
    "        plt.cla()\n",
    "        plt.close()\n",
    "        plt.close(fig1)\n",
    "\n",
    "        fig2 = plt.figure(figsize = figsize)\n",
    "        plt.axes(ylim=(-1, 1))\n",
    "        sigr = signal.real\n",
    "        sigi = signal.imag\n",
    "        \n",
    "        plt.plot(sigi, color='black')\n",
    "        plt.gca().set_axis_off()\n",
    "        plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)\n",
    "        plt.margins(0,0)\n",
    "        buf = io.BytesIO()\n",
    "        fig2.savefig(buf, format=\"png\", dpi=dpi)\n",
    "        buf.seek(0)\n",
    "        img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)\n",
    "        buf.close()\n",
    "        img = cv2.imdecode(img_arr, 1)\n",
    "        img2 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "        plt.clf()\n",
    "        plt.cla()\n",
    "        plt.close()\n",
    "        plt.close(fig2)\n",
    "\n",
    "        img = np.array([img1, img2, spectr])\n",
    " \n",
    "        cv2.destroyAllWindows()\n",
    "        del signal\n",
    "        del spec\n",
    "        del spectr\n",
    "        del img1\n",
    "        del img2\n",
    "        del sigr\n",
    "        del sigi\n",
    "        del buf\n",
    "        del buf1\n",
    "        del img_arr\n",
    "        del img_arr1\n",
    "        cv2.destroyAllWindows()\n",
    "        gc.collect()\n",
    "\n",
    "        print('Подготовка данных завершена')\n",
    "        print()\n",
    "        return img\n",
    "\n",
    "    except Exception as e:\n",
    "        print(str(e))\n",
    "        return None\n",
    "\n",
    "\n",
    "def build_func_ensemble(file_model='', file_config='', num_classes=None):\n",
    "    try:\n",
    "        import matplotlib.pyplot as plt\n",
    "        matplotlib.use('Agg')\n",
    "        plt.ioff()\n",
    "        torch.cuda.empty_cache()\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)\n",
    "                x2 = self.model2(x)\n",
    "                x3 = self.model3(x)\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",
    "        device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "        if device != 'cpu':\n",
    "            model = model.to(device)\n",
    "        model.load_state_dict(torch.load(file_model, map_location=device))\n",
    "        model.eval()\n",
    "\n",
    "        cv2.destroyAllWindows()\n",
    "        del model1\n",
    "        del model2\n",
    "        del model3\n",
    "        gc.collect()\n",
    "\n",
    "        print('Инициализация модели завершена')\n",
    "        print()\n",
    "        return model\n",
    "\n",
    "    except Exception as exc:\n",
    "        print(str(exc))\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "448da74a-e0ae-44d8-9877-8dd1f257a24f",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_binaries = 'C:/Users/snytk/Lerning_NN_for_work/datasets/1.2_learning'\n",
    "path_to_pictures = 'C:/Users/snytk/Lerning_NN_for_work/datasets_jpg/1.2_jpg_learning'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ac4945a8-29c4-4da4-945f-08658953e3e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f226f86-5d72-4573-8af6-750128b70263",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 963/963 [06:36<00:00,  2.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dir:  drone  finished!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████| 1724/1724 [11:41<00:00,  2.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dir:  noise  finished!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "size = (256,256)\n",
    "if not os.path.exists(path_to_pictures):\n",
    "    os.mkdir(path_to_pictures)\n",
    "for subdir in os.listdir(path_to_binaries):\n",
    "    filepath = path_to_binaries + '/' + subdir + '/'\n",
    "    if not os.path.exists(path_to_pictures +'/' + subdir):\n",
    "        os.mkdir(path_to_pictures + '/' + subdir)\n",
    "    files = os.listdir(filepath)\n",
    "    for file in tqdm(files):\n",
    "        savepath_real_png = path_to_pictures +'/' + subdir + '/' + file + '_real' + '.jpg' \n",
    "        savepath_imag_png = path_to_pictures +'/' + subdir + '/' + file + '_imag' + '.jpg' \n",
    "        savepath_spec_png = path_to_pictures +'/' + subdir + '/' + file + '_spec' + '.jpg'\n",
    "        if not os.path.exists(savepath_real_png):\n",
    "            img = sig2pic_with_spec(path_to_data=filepath, filename=file, resize = size)\n",
    "            gc.collect()\n",
    "            \n",
    "            try:\n",
    "                plt.imshow(img[0])\n",
    "                plt.savefig(savepath_real_png)\n",
    "                plt.clf()\n",
    "                plt.cla()\n",
    "                plt.close()\n",
    "                \n",
    "                plt.imshow(img[1])\n",
    "                plt.savefig(savepath_imag_png)\n",
    "                plt.clf()\n",
    "                plt.cla()\n",
    "                plt.close()\n",
    "\n",
    "                plt.imshow(img[2])\n",
    "                plt.savefig(savepath_spec_png)\n",
    "                plt.clf()\n",
    "                plt.cla()\n",
    "                plt.close()\n",
    "                \n",
    "            except Exception:\n",
    "                continue\n",
    "    print('Dir: ', subdir , ' finished!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "871d7ab6",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
