{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Шаблоны сохранены в папке 'templates'.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "\n",
    "# Убедитесь, что у вас установлен шрифт RoadNumbers2.0.ttf\n",
    "FONT_PATH = \"RoadNumbers2.0.ttf\"\n",
    "OUTPUT_DIR = \"templates\"\n",
    "\n",
    "# Буквы и цифры, используемые в российских номерах\n",
    "SYMBOLS = \"DABEKMHOPCTYX0123456789\"\n",
    "IMG_SIZE = (28, 28)\n",
    "FONT_SIZE = 24\n",
    "\n",
    "# Создаем директорию для шаблонов\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# Загружаем шрифт\n",
    "try:\n",
    "    font = ImageFont.truetype(FONT_PATH, FONT_SIZE)\n",
    "except IOError:\n",
    "    print(f\"Шрифт {FONT_PATH} не найден. Убедитесь, что он находится в текущей директории.\")\n",
    "    exit()\n",
    "\n",
    "# Генерация изображений для каждого символа\n",
    "for symbol in SYMBOLS:\n",
    "    # Создаем пустое изображение с белым фоном\n",
    "    img = Image.new(\"L\", IMG_SIZE, 255)\n",
    "    draw = ImageDraw.Draw(img)\n",
    "\n",
    "    # Получаем размеры текста с использованием Font.getbbox\n",
    "    text_width, text_height = draw.textbbox((0, 0), symbol, font=font)[2:] # [2:] берёт ширину и высоту\n",
    "\n",
    "    # Рассчитываем координаты для центрирования текста\n",
    "    x = (IMG_SIZE[0] - text_width) // 2\n",
    "    y = (IMG_SIZE[1] - text_height) // 2\n",
    "\n",
    "    # Рисуем текст\n",
    "    draw.text((x, y), symbol, font=font, fill=0)\n",
    "\n",
    "    # Сохраняем изображение\n",
    "    img.save(os.path.join(OUTPUT_DIR, f\"{symbol}.png\"))\n",
    "\n",
    "print(f\"Шаблоны сохранены в папке '{OUTPUT_DIR}'.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 153\u001b[0m\n\u001b[0;32m    150\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mКонфигурационный файл создан: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mCONFIG_FILE\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    152\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m--> 153\u001b[0m     \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[2], line 142\u001b[0m, in \u001b[0;36mmain\u001b[1;34m()\u001b[0m\n\u001b[0;32m    139\u001b[0m config_lines \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m    141\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m symbol \u001b[38;5;129;01min\u001b[39;00m SYMBOLS:\n\u001b[1;32m--> 142\u001b[0m     \u001b[43mcreate_augmentations\u001b[49m\u001b[43m(\u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mFONT_SIZES\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mMAX_ROTATION\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mAUGMENTATIONS\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig_lines\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    144\u001b[0m \u001b[38;5;66;03m# Записываем все строки конфигурации одним разом\u001b[39;00m\n\u001b[0;32m    145\u001b[0m config_path \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(OUTPUT_DIR, CONFIG_FILE)\n",
      "Cell \u001b[1;32mIn[2], line 74\u001b[0m, in \u001b[0;36mcreate_augmentations\u001b[1;34m(symbol, font_sizes, max_rotation, augmentations, config_lines)\u001b[0m\n\u001b[0;32m     72\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(augmentations):\n\u001b[0;32m     73\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m font_size \u001b[38;5;129;01min\u001b[39;00m font_sizes:\n\u001b[1;32m---> 74\u001b[0m         current_font \u001b[38;5;241m=\u001b[39m \u001b[43mImageFont\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtruetype\u001b[49m\u001b[43m(\u001b[49m\u001b[43mFONT_PATH\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfont_size\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     76\u001b[0m         \u001b[38;5;66;03m# Создаём пустое изображение с белым фоном\u001b[39;00m\n\u001b[0;32m     77\u001b[0m         img \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mnew(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mL\u001b[39m\u001b[38;5;124m\"\u001b[39m, IMG_SIZE, \u001b[38;5;241m255\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\leonk\\Documents\\code\\study\\.venv\\lib\\site-packages\\PIL\\ImageFont.py:879\u001b[0m, in \u001b[0;36mtruetype\u001b[1;34m(font, size, index, encoding, layout_engine)\u001b[0m\n\u001b[0;32m    876\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m FreeTypeFont(font, size, index, encoding, layout_engine)\n\u001b[0;32m    878\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 879\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfreetype\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfont\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    880\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m:\n\u001b[0;32m    881\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_path(font):\n",
      "File \u001b[1;32mc:\\Users\\leonk\\Documents\\code\\study\\.venv\\lib\\site-packages\\PIL\\ImageFont.py:876\u001b[0m, in \u001b[0;36mtruetype.<locals>.freetype\u001b[1;34m(font)\u001b[0m\n\u001b[0;32m    875\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfreetype\u001b[39m(font: StrOrBytesPath \u001b[38;5;241m|\u001b[39m BinaryIO) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m FreeTypeFont:\n\u001b[1;32m--> 876\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mFreeTypeFont\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfont\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msize\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mencoding\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlayout_engine\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\leonk\\Documents\\code\\study\\.venv\\lib\\site-packages\\PIL\\ImageFont.py:273\u001b[0m, in \u001b[0;36mFreeTypeFont.__init__\u001b[1;34m(self, font, size, index, encoding, layout_engine)\u001b[0m\n\u001b[0;32m    268\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfont \u001b[38;5;241m=\u001b[39m core\u001b[38;5;241m.\u001b[39mgetfont(\n\u001b[0;32m    269\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m, size, index, encoding, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfont_bytes, layout_engine\n\u001b[0;32m    270\u001b[0m     )\n\u001b[0;32m    272\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m is_path(font):\n\u001b[1;32m--> 273\u001b[0m     font \u001b[38;5;241m=\u001b[39m \u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrealpath\u001b[49m\u001b[43m(\u001b[49m\u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfspath\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfont\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    274\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m sys\u001b[38;5;241m.\u001b[39mplatform \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwin32\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m    275\u001b[0m         font_bytes_path \u001b[38;5;241m=\u001b[39m font \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(font, \u001b[38;5;28mbytes\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m font\u001b[38;5;241m.\u001b[39mencode()\n",
      "File \u001b[1;32m~\\AppData\\Local\\Programs\\Python\\Python39\\lib\\ntpath.py:664\u001b[0m, in \u001b[0;36mrealpath\u001b[1;34m(path)\u001b[0m\n\u001b[0;32m    662\u001b[0m \u001b[38;5;66;03m# Ensure that the non-prefixed path resolves to the same path\u001b[39;00m\n\u001b[0;32m    663\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 664\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[43m_getfinalpathname\u001b[49m\u001b[43m(\u001b[49m\u001b[43mspath\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m==\u001b[39m path:\n\u001b[0;32m    665\u001b[0m         path \u001b[38;5;241m=\u001b[39m spath\n\u001b[0;32m    666\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m ex:\n\u001b[0;32m    667\u001b[0m     \u001b[38;5;66;03m# If the path does not exist and originally did not exist, then\u001b[39;00m\n\u001b[0;32m    668\u001b[0m     \u001b[38;5;66;03m# strip the prefix anyway.\u001b[39;00m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageFilter\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "# Убедитесь, что у вас установлен шрифт RoadNumbers2.0.ttf\n",
    "FONT_PATH = \"RoadNumbers2.0.ttf\"\n",
    "OUTPUT_DIR = \"dataset\"\n",
    "CONFIG_FILE = \"dataset_config.txt\"\n",
    "\n",
    "# Буквы и цифры, используемые в российских номерах\n",
    "SYMBOLS = \"ABEKMHOPCTYX0123456789\"\n",
    "IMG_SIZE = (28, 28)\n",
    "FONT_SIZES = [26, 32, 38, 46, 54, 58]\n",
    "MAX_ROTATION = 15\n",
    "AUGMENTATIONS = 160\n",
    "\n",
    "# Создаём директорию для датасета\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# Загружаем шрифт\n",
    "try:\n",
    "    font = ImageFont.truetype(FONT_PATH, FONT_SIZES[0])\n",
    "except IOError:\n",
    "    print(f\"Шрифт {FONT_PATH} не найден. Убедитесь, что он находится в текущей директории.\")\n",
    "    exit()\n",
    "\n",
    "# Функция расширения объектов\n",
    "def expand_characters(img, kernel_size=(2, 2), iterations=1):\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)\n",
    "    img_array = np.array(img)\n",
    "    expanded = cv2.dilate(img_array, kernel, iterations=iterations)\n",
    "    return Image.fromarray(expanded)\n",
    "\n",
    "# Функция для добавления случайных разрывов в контуре\n",
    "def add_random_gaps(img, num_gaps=5, gap_size=5):\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    for _ in range(num_gaps):\n",
    "        x1 = random.randint(0, IMG_SIZE[0] - gap_size)\n",
    "        y1 = random.randint(0, IMG_SIZE[1] - gap_size)\n",
    "        x2 = x1 + gap_size\n",
    "        y2 = y1 + gap_size\n",
    "        draw.rectangle([x1, y1, x2, y2], fill=255)  # Белый прямоугольник для разрыва\n",
    "    return img\n",
    "\n",
    "# Функция для добавления блюра\n",
    "def apply_blur(img):\n",
    "    blur_type = random.choice([\"box\", \"gaussian\", \"motion\"])\n",
    "    if blur_type == \"box\":\n",
    "        img = img.filter(ImageFilter.BoxBlur(random.randint(1, 4)))\n",
    "    elif blur_type == \"gaussian\":\n",
    "        img = img.filter(ImageFilter.GaussianBlur(random.randint(1, 3)))\n",
    "    elif blur_type == \"motion\":\n",
    "        kernel_size = random.randint(4, 8)\n",
    "        kernel_motion_blur = np.zeros((kernel_size, kernel_size))\n",
    "        kernel_motion_blur[int((kernel_size - 1) / 2), :] = 1\n",
    "        kernel_motion_blur /= kernel_size\n",
    "        img_array = cv2.filter2D(np.array(img), -1, kernel_motion_blur)\n",
    "        img = Image.fromarray(img_array)\n",
    "    return img\n",
    "\n",
    "# Функция для добавления шума\n",
    "def add_noise(img, intensity=30):\n",
    "    img_array = np.array(img)\n",
    "    noise = np.random.normal(0, intensity, img_array.shape).astype(np.int32)\n",
    "    noisy_img = np.clip(img_array + noise, 0, 255).astype(np.uint8)\n",
    "    return Image.fromarray(noisy_img)\n",
    "\n",
    "# Функция для создания аугментаций\n",
    "def create_augmentations(symbol, font_sizes, max_rotation, augmentations, config_lines):\n",
    "    for i in range(augmentations):\n",
    "        for font_size in font_sizes:\n",
    "            current_font = ImageFont.truetype(FONT_PATH, font_size)\n",
    "\n",
    "            # Создаём пустое изображение с белым фоном\n",
    "            img = Image.new(\"L\", IMG_SIZE, 255)\n",
    "            draw = ImageDraw.Draw(img)\n",
    "\n",
    "            # Получаем размеры текста с использованием Font.getbbox\n",
    "            text_bbox = draw.textbbox((0, 0), symbol, font=current_font)\n",
    "            text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]\n",
    "\n",
    "            # Рассчитываем координаты для центрирования текста\n",
    "            x = (IMG_SIZE[0] - text_width) // 2 - text_bbox[0]\n",
    "            y = (IMG_SIZE[1] - text_height) // 2 - text_bbox[1]\n",
    "\n",
    "            # Рисуем текст\n",
    "            draw.text((x, y), symbol, font=current_font, fill=0)\n",
    "\n",
    "            # Для первого прогона символа с данным размером шрифта сохраняем без аугментаций\n",
    "            if i == 0:\n",
    "                symbol_dir = os.path.join(OUTPUT_DIR, symbol)\n",
    "                os.makedirs(symbol_dir, exist_ok=True)\n",
    "                file_path = os.path.join(symbol_dir, f\"{symbol}_{font_size}_{i}.png\")\n",
    "                img.save(file_path)\n",
    "\n",
    "                config_lines.append(f\"{file_path},{symbol}\\n\")\n",
    "                continue\n",
    "\n",
    "            # Случайный поворот изображения\n",
    "            rotation = random.uniform(-max_rotation, max_rotation)\n",
    "            img = img.rotate(rotation, expand=False, fillcolor=255)\n",
    "\n",
    "            # Случайное расширение символов\n",
    "            if random.random() < 0.5:\n",
    "                kernel_size = (random.randint(1, 2), random.randint(1, 2))\n",
    "                iterations = 1\n",
    "                img = expand_characters(img, kernel_size=kernel_size, iterations=iterations)\n",
    "\n",
    "            # Случайное добавление разрывов\n",
    "            if random.random() < 0.75:\n",
    "                num_gaps = random.randint(1, 5)\n",
    "                gap_size = random.randint(2, 4)\n",
    "                img = add_random_gaps(img, num_gaps=num_gaps, gap_size=gap_size)\n",
    "\n",
    "            if random.random() < 0.25:\n",
    "                img = ImageOps.invert(img)\n",
    "\n",
    "            # Случайное добавление блюра\n",
    "            if random.random() < 0.85:\n",
    "                img = apply_blur(img)\n",
    "\n",
    "            # Случайное добавление шума\n",
    "            if random.random() < 0.1:\n",
    "                noise_intensity = random.randint(10, 20)\n",
    "                img = add_noise(img, intensity=noise_intensity)\n",
    "\n",
    "            # Сохраняем изображение\n",
    "            symbol_dir = os.path.join(OUTPUT_DIR, symbol)\n",
    "            os.makedirs(symbol_dir, exist_ok=True)\n",
    "            file_path = os.path.join(symbol_dir, f\"{symbol}_{font_size}_{i}.png\")\n",
    "            img.save(file_path)\n",
    "\n",
    "            config_lines.append(f\"{file_path},{symbol}\\n\")\n",
    "\n",
    "# Основной блок генерации датасета\n",
    "def main():\n",
    "    config_lines = []\n",
    "\n",
    "    for symbol in SYMBOLS:\n",
    "        create_augmentations(symbol, FONT_SIZES, MAX_ROTATION, AUGMENTATIONS, config_lines)\n",
    "\n",
    "    # Записываем все строки конфигурации одним разом\n",
    "    config_path = os.path.join(OUTPUT_DIR, CONFIG_FILE)\n",
    "    with open(config_path, \"w\") as config:\n",
    "        config.writelines(config_lines)\n",
    "\n",
    "    print(f\"Датасет сохранен в папке '{OUTPUT_DIR}'.\")\n",
    "    print(f\"Конфигурационный файл создан: {CONFIG_FILE}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Датасет сохранен в папке 'dataset'.\n",
      "Конфигурационный файл создан: dataset_config.txt\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageFilter\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "# Убедитесь, что у вас установлен шрифт RoadNumbers2.0.ttf\n",
    "FONT_PATH = \"RoadNumbers2.0.ttf\"\n",
    "OUTPUT_DIR = \"dataset\"\n",
    "CONFIG_FILE = \"dataset_config.txt\"\n",
    "\n",
    "# Буквы и цифры, используемые в российских номерах\n",
    "SYMBOLS = \"ABEKMHOPCTYX0123456789\"\n",
    "IMG_SIZE = (28, 28)\n",
    "FONT_SIZES = [26, 32, 38, 46, 54, 58]\n",
    "MAX_ROTATION = 15\n",
    "AUGMENTATIONS = 160\n",
    "\n",
    "# Создаём директорию для датасета\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# Загружаем шрифт\n",
    "try:\n",
    "    font = ImageFont.truetype(FONT_PATH, FONT_SIZES[0])\n",
    "except IOError:\n",
    "    print(f\"Шрифт {FONT_PATH} не найден. Убедитесь, что он находится в текущей директории.\")\n",
    "    exit()\n",
    "\n",
    "# Функция расширения объектов\n",
    "def expand_characters(img, kernel_size=(2, 2), iterations=1):\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)\n",
    "    img_array = np.array(img)\n",
    "    expanded = cv2.dilate(img_array, kernel, iterations=iterations)\n",
    "    return Image.fromarray(expanded)\n",
    "\n",
    "# Функция для добавления случайных разрывов в контуре\n",
    "def add_random_gaps(img, num_gaps=5, gap_size=5):\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    for _ in range(num_gaps):\n",
    "        x1 = random.randint(0, IMG_SIZE[0] - gap_size)\n",
    "        y1 = random.randint(0, IMG_SIZE[1] - gap_size)\n",
    "        x2 = x1 + gap_size\n",
    "        y2 = y1 + gap_size\n",
    "        draw.rectangle([x1, y1, x2, y2], fill=255)  # Белый прямоугольник для разрыва\n",
    "    return img\n",
    "\n",
    "# Функция для добавления блюра\n",
    "def apply_blur(img):\n",
    "    blur_type = random.choice([\"box\", \"gaussian\", \"motion\"])\n",
    "    if blur_type == \"box\":\n",
    "        img = img.filter(ImageFilter.BoxBlur(random.randint(1, 4)))\n",
    "    elif blur_type == \"gaussian\":\n",
    "        img = img.filter(ImageFilter.GaussianBlur(random.randint(1, 3)))\n",
    "    elif blur_type == \"motion\":\n",
    "        kernel_size = random.randint(4, 8)\n",
    "        kernel_motion_blur = np.zeros((kernel_size, kernel_size))\n",
    "        kernel_motion_blur[int((kernel_size - 1) / 2), :] = 1\n",
    "        kernel_motion_blur /= kernel_size\n",
    "        img_array = cv2.filter2D(np.array(img), -1, kernel_motion_blur)\n",
    "        img = Image.fromarray(img_array)\n",
    "    return img\n",
    "\n",
    "# Функция для добавления шума\n",
    "def add_noise(img, intensity=30):\n",
    "    img_array = np.array(img)\n",
    "    noise = np.random.normal(0, intensity, img_array.shape).astype(np.int32)\n",
    "    noisy_img = np.clip(img_array + noise, 0, 255).astype(np.uint8)\n",
    "    return Image.fromarray(noisy_img)\n",
    "\n",
    "# Новая функция для смещения изображения\n",
    "def shift_image(img, max_shift=4):\n",
    "    \"\"\"\n",
    "    Сдвигает изображение на случайное количество пикселей по осям X и Y.\n",
    "\n",
    "    :param img: PIL Image объект.\n",
    "    :param max_shift: Максимальное смещение в пикселях по каждой оси.\n",
    "    :return: Сдвинутое изображение.\n",
    "    \"\"\"\n",
    "    x_shift = random.randint(-max_shift, max_shift)\n",
    "    y_shift = random.randint(-max_shift, max_shift)\n",
    "    return img.transform(\n",
    "        img.size,\n",
    "        Image.AFFINE,\n",
    "        (1, 0, x_shift, 0, 1, y_shift),\n",
    "        fillcolor=255  # Заполнение пустых областей белым цветом\n",
    "    )\n",
    "\n",
    "# Функция для создания аугментаций\n",
    "def create_augmentations(symbol, font_sizes, max_rotation, augmentations, config_lines):\n",
    "    for i in range(augmentations):\n",
    "        for font_size in font_sizes:\n",
    "            current_font = ImageFont.truetype(FONT_PATH, font_size)\n",
    "\n",
    "            # Создаём пустое изображение с белым фоном\n",
    "            img = Image.new(\"L\", IMG_SIZE, 255)\n",
    "            draw = ImageDraw.Draw(img)\n",
    "\n",
    "            # Получаем размеры текста с использованием Font.getbbox\n",
    "            text_bbox = draw.textbbox((0, 0), symbol, font=current_font)\n",
    "            text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]\n",
    "\n",
    "            # Рассчитываем координаты для центрирования текста\n",
    "            x = (IMG_SIZE[0] - text_width) // 2 - text_bbox[0]\n",
    "            y = (IMG_SIZE[1] - text_height) // 2 - text_bbox[1]\n",
    "\n",
    "            # Рисуем текст\n",
    "            draw.text((x, y), symbol, font=current_font, fill=0)\n",
    "\n",
    "            # Для первого прогона символа с данным размером шрифта сохраняем без аугментаций\n",
    "            if i == 0:\n",
    "                symbol_dir = os.path.join(OUTPUT_DIR, symbol)\n",
    "                os.makedirs(symbol_dir, exist_ok=True)\n",
    "                file_path = os.path.join(symbol_dir, f\"{symbol}_{font_size}_{i}.png\")\n",
    "                img.save(file_path)\n",
    "\n",
    "                config_lines.append(f\"{file_path},{symbol}\\n\")\n",
    "                continue\n",
    "\n",
    "            # Случайное смещение символа\n",
    "            if random.random() < 0.7:  # 50% вероятность применения смещения\n",
    "                img = shift_image(img, max_shift=4)\n",
    "\n",
    "            # Случайный поворот изображения\n",
    "            rotation = random.uniform(-max_rotation, max_rotation)\n",
    "            img = img.rotate(rotation, expand=False, fillcolor=255)\n",
    "\n",
    "            # Случайное расширение символов\n",
    "            if random.random() < 0.5:\n",
    "                kernel_size = (random.randint(1, 2), random.randint(1, 2))\n",
    "                iterations = 1\n",
    "                img = expand_characters(img, kernel_size=kernel_size, iterations=iterations)\n",
    "\n",
    "            # Случайное добавление разрывов\n",
    "            if random.random() < 0.6:\n",
    "                num_gaps = random.randint(1, 5)\n",
    "                gap_size = random.randint(2, 4)\n",
    "                img = add_random_gaps(img, num_gaps=num_gaps, gap_size=gap_size)\n",
    "\n",
    "            # Инверсия цветов с вероятностью 25%\n",
    "            if random.random() < 0.15:\n",
    "                img = ImageOps.invert(img)\n",
    "\n",
    "            # Случайное добавление блюра\n",
    "            if random.random() < 0.85:\n",
    "                img = apply_blur(img)\n",
    "\n",
    "            # Случайное добавление шума\n",
    "            if random.random() < 0.1:\n",
    "                noise_intensity = random.randint(10, 20)\n",
    "                img = add_noise(img, intensity=noise_intensity)\n",
    "\n",
    "            # Сохраняем изображение\n",
    "            symbol_dir = os.path.join(OUTPUT_DIR, symbol)\n",
    "            os.makedirs(symbol_dir, exist_ok=True)\n",
    "            file_path = os.path.join(symbol_dir, f\"{symbol}_{font_size}_{i}.png\")\n",
    "            img.save(file_path)\n",
    "\n",
    "            config_lines.append(f\"{file_path},{symbol}\\n\")\n",
    "\n",
    "# Основной блок генерации датасета\n",
    "def main():\n",
    "    config_lines = []\n",
    "\n",
    "    for symbol in SYMBOLS:\n",
    "        create_augmentations(symbol, FONT_SIZES, MAX_ROTATION, AUGMENTATIONS, config_lines)\n",
    "\n",
    "    # Записываем все строки конфигурации одним разом\n",
    "    config_path = os.path.join(OUTPUT_DIR, CONFIG_FILE)\n",
    "    with open(config_path, \"w\") as config:\n",
    "        config.writelines(config_lines)\n",
    "\n",
    "    print(f\"Датасет сохранен в папке '{OUTPUT_DIR}'.\")\n",
    "    print(f\"Конфигурационный файл создан: {CONFIG_FILE}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18100 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import cv2\n",
    "from PIL import Image, ImageDraw, ImageFont, ImageFilter\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Путь к шрифту\n",
    "FONT_PATH = \"RoadNumbers2.0.ttf\"  # Убедитесь, что шрифт находится в текущей директории\n",
    "FONT_SIZE = 26  # Размер шрифта (можно изменить при необходимости)\n",
    "IMG_SIZE = (28, 28)  # Размер изображения\n",
    "\n",
    "# Функция для создания изображения с символом\n",
    "def create_symbol_image(symbol, font_path, font_size, img_size):\n",
    "    try:\n",
    "        font = ImageFont.truetype(font_path, font_size)\n",
    "    except IOError:\n",
    "        print(f\"Шрифт {font_path} не найден. Убедитесь, что он находится в текущей директории.\")\n",
    "        exit()\n",
    "    \n",
    "    # Создаём пустое изображение с белым фоном\n",
    "    img = Image.new(\"L\", img_size, 255)  # \"L\" для градаций серого\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    \n",
    "    # Получаем размеры текста\n",
    "    text_bbox = draw.textbbox((0, 0), symbol, font=font)\n",
    "    text_width = text_bbox[2] - text_bbox[0]\n",
    "    text_height = text_bbox[3] - text_bbox[1]\n",
    "    \n",
    "    # Рассчитываем координаты для центрирования текста\n",
    "    x = (img_size[0] - text_width) // 2 - text_bbox[0]\n",
    "    y = (img_size[1] - text_height) // 2 - text_bbox[1]\n",
    "    \n",
    "    # Рисуем текст\n",
    "    draw.text((x, y), symbol, font=font, fill=0)  # Черный цвет\n",
    "    \n",
    "    return img\n",
    "\n",
    "# Функция для применения блюра\n",
    "def apply_blur(img, blur_type):\n",
    "    if blur_type == \"box\":\n",
    "        radius = 4\n",
    "        blurred_img = img.filter(ImageFilter.BoxBlur(radius))\n",
    "    elif blur_type == \"gaussian\":\n",
    "        radius = 3\n",
    "        blurred_img = img.filter(ImageFilter.GaussianBlur(radius))\n",
    "    elif blur_type == \"motion\":\n",
    "        kernel_size = 8\n",
    "        # Создаём ядро для motion blur (размытие по горизонтали)\n",
    "        kernel_motion_blur = np.zeros((kernel_size, kernel_size))\n",
    "        kernel_motion_blur[int((kernel_size - 1) / 2), :] = 1\n",
    "        kernel_motion_blur /= kernel_size\n",
    "        img_array = np.array(img)\n",
    "        blurred_array = cv2.filter2D(img_array, -1, kernel_motion_blur)\n",
    "        blurred_img = Image.fromarray(blurred_array)\n",
    "    else:\n",
    "        raise ValueError(\"Неподдерживаемый тип блюра.\")\n",
    "    \n",
    "    return blurred_img\n",
    "\n",
    "# Основная функция\n",
    "def main():\n",
    "    symbol = '7'  # Цифра для обработки\n",
    "    blur_types = [\"box\", \"gaussian\", \"motion\"]\n",
    "    \n",
    "    # Создаём оригинальное изображение\n",
    "    original_img = create_symbol_image(symbol, FONT_PATH, FONT_SIZE, IMG_SIZE)\n",
    "    \n",
    "    # Применяем каждый тип блюра\n",
    "    blurred_images = {}\n",
    "    for blur_type in blur_types:\n",
    "        blurred_img = apply_blur(original_img, blur_type)\n",
    "        blurred_images[blur_type] = blurred_img\n",
    "    \n",
    "    # Настройка отображения с помощью matplotlib\n",
    "    num_blurs = len(blur_types)\n",
    "    plt.figure(figsize=(4 * (num_blurs + 1), 4))  # Размер фигуры зависит от количества изображений\n",
    "    \n",
    "    # Отображаем оригинальное изображение\n",
    "    plt.subplot(1, num_blurs + 1, 1)\n",
    "    plt.imshow(original_img, cmap='gray')\n",
    "    plt.title(\"Оригинал\")\n",
    "    plt.axis('off')\n",
    "    \n",
    "    # Отображаем размытые изображения\n",
    "    for idx, blur_type in enumerate(blur_types, start=2):\n",
    "        plt.subplot(1, num_blurs + 1, idx)\n",
    "        plt.imshow(blurred_images[blur_type], cmap='gray')\n",
    "        plt.title(f\"{blur_type.capitalize()} Blur\")\n",
    "        plt.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "DATASET_DIR = \"dataset\"\n",
    "SAVE_PATH = \"best_model_3.pth\"\n",
    "BATCH_SIZE = 32\n",
    "EPOCHS = 30\n",
    "LEARNING_RATE = 0.001\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "CLASSES = \"ABEKMHOPCTYX0123456789\"\n",
    "NUM_CLASSES = len(CLASSES)\n",
    "CLASS_TO_IDX = {char: idx for idx, char in enumerate(CLASSES)}\n",
    "IDX_TO_CLASS = {idx: char for char, idx in CLASS_TO_IDX.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Размер тренировочного набора: 16896\n",
      "Размер тестового набора: 4224\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "class LicensePlateDataset(Dataset):\n",
    "    def __init__(self, dataset_dir, transform=None):\n",
    "        self.data = []\n",
    "        self.transform = transform\n",
    "        with open(os.path.join(dataset_dir, \"dataset_config.txt\"), \"r\") as file:\n",
    "            for line in file:\n",
    "                path, label = line.strip().split(\",\")\n",
    "                self.data.append((path, CLASS_TO_IDX[label]))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path, label = self.data[idx]\n",
    "        image = Image.open(img_path).convert(\"L\")\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label\n",
    "\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ColorJitter(brightness=0.4, contrast=0.4),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.5], std=[0.5]),\n",
    "])\n",
    "\n",
    "full_dataset = LicensePlateDataset(DATASET_DIR, transform=transform)\n",
    "\n",
    "indices = list(range(len(full_dataset)))\n",
    "train_indices, test_indices = train_test_split(indices, test_size=0.2, random_state=42)\n",
    "\n",
    "train_dataset = Subset(full_dataset, train_indices)\n",
    "test_dataset = Subset(full_dataset, test_indices)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "print(f\"Размер тренировочного набора: {len(train_dataset)}\")\n",
    "print(f\"Размер тестового набора: {len(test_dataset)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),  # Увеличиваем количество фильтров\n",
    "            nn.LeakyReLU(negative_slope=0.1),  # Используем LeakyReLU вместо ReLU\n",
    "            nn.BatchNorm2d(32),  # Нормализация\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            \n",
    "            nn.Conv2d(32, 64, kernel_size=3, padding=1),\n",
    "            nn.ELU(),  # Используем ELU вместо ReLU\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            \n",
    "            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n",
    "            nn.SiLU(),  # Используем SiLU (Swish) вместо ReLU\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Linear(128 * 3 * 3, 256),\n",
    "            nn.GELU(),  # Используем GELU для полносвязного слоя\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(256, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc_layers(x)\n",
    "        return x\n",
    "\n",
    "model = SimpleCNN(len(CLASSES)).to(DEVICE)\n",
    "criterion = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=1e-4)\n",
    "#optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)\n",
    "\n",
    "#optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training:   0%|          | 0/528 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 2.2562, Train Accuracy: 41.64%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.5400, Test Accuracy: 72.47%\n",
      "Best model saved with test accuracy: 72.47%\n",
      "Epoch [2/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.4775, Train Accuracy: 72.54%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.2487, Test Accuracy: 82.43%\n",
      "Best model saved with test accuracy: 82.43%\n",
      "Epoch [3/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.2714, Train Accuracy: 80.63%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.1178, Test Accuracy: 86.65%\n",
      "Best model saved with test accuracy: 86.65%\n",
      "Epoch [4/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.1700, Train Accuracy: 84.59%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.0693, Test Accuracy: 88.02%\n",
      "Best model saved with test accuracy: 88.02%\n",
      "Epoch [5/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.0998, Train Accuracy: 87.51%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.0128, Test Accuracy: 89.89%\n",
      "Best model saved with test accuracy: 89.89%\n",
      "Epoch [6/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.0531, Train Accuracy: 89.01%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.9741, Test Accuracy: 90.77%\n",
      "Best model saved with test accuracy: 90.77%\n",
      "Epoch [7/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.0165, Train Accuracy: 90.26%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.9434, Test Accuracy: 92.14%\n",
      "Best model saved with test accuracy: 92.14%\n",
      "Epoch [8/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9889, Train Accuracy: 91.57%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.9299, Test Accuracy: 93.42%\n",
      "Best model saved with test accuracy: 93.42%\n",
      "Epoch [9/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9668, Train Accuracy: 92.20%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.9245, Test Accuracy: 92.57%\n",
      "Epoch [10/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9427, Train Accuracy: 92.90%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.9001, Test Accuracy: 92.87%\n",
      "Epoch [11/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9336, Train Accuracy: 92.98%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8980, Test Accuracy: 93.11%\n",
      "Epoch [12/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9161, Train Accuracy: 93.77%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8720, Test Accuracy: 94.34%\n",
      "Best model saved with test accuracy: 94.34%\n",
      "Epoch [13/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.9031, Train Accuracy: 94.06%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8702, Test Accuracy: 93.92%\n",
      "Epoch [14/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8949, Train Accuracy: 94.36%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8564, Test Accuracy: 94.79%\n",
      "Best model saved with test accuracy: 94.79%\n",
      "Epoch [15/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8834, Train Accuracy: 94.86%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8516, Test Accuracy: 94.77%\n",
      "Epoch [16/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8761, Train Accuracy: 94.99%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8566, Test Accuracy: 94.11%\n",
      "Epoch [17/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8690, Train Accuracy: 95.05%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8463, Test Accuracy: 94.86%\n",
      "Best model saved with test accuracy: 94.86%\n",
      "Epoch [18/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8607, Train Accuracy: 95.47%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8430, Test Accuracy: 94.77%\n",
      "Epoch [19/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8550, Train Accuracy: 95.57%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8356, Test Accuracy: 94.84%\n",
      "Epoch [20/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8473, Train Accuracy: 95.70%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8284, Test Accuracy: 94.72%\n",
      "Epoch [21/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8450, Train Accuracy: 95.86%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8287, Test Accuracy: 95.08%\n",
      "Best model saved with test accuracy: 95.08%\n",
      "Epoch [22/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8378, Train Accuracy: 96.04%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8208, Test Accuracy: 95.36%\n",
      "Best model saved with test accuracy: 95.36%\n",
      "Epoch [23/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8315, Train Accuracy: 96.20%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8261, Test Accuracy: 95.31%\n",
      "Epoch [24/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8314, Train Accuracy: 96.29%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8259, Test Accuracy: 95.38%\n",
      "Best model saved with test accuracy: 95.38%\n",
      "Epoch [25/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8240, Train Accuracy: 96.51%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8189, Test Accuracy: 94.91%\n",
      "Epoch [26/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8187, Train Accuracy: 96.73%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8071, Test Accuracy: 95.41%\n",
      "Best model saved with test accuracy: 95.41%\n",
      "Epoch [27/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8145, Train Accuracy: 96.63%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8207, Test Accuracy: 95.27%\n",
      "Epoch [28/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8134, Train Accuracy: 96.76%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8146, Test Accuracy: 95.50%\n",
      "Best model saved with test accuracy: 95.50%\n",
      "Epoch [29/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8097, Train Accuracy: 96.77%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8127, Test Accuracy: 95.12%\n",
      "Epoch [30/30]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.8054, Train Accuracy: 96.92%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.8038, Test Accuracy: 95.55%\n",
      "Best model saved with test accuracy: 95.55%\n",
      "Обучение завершено.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "best_accuracy = 0.0\n",
    "for epoch in range(EPOCHS):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    correct_train = 0\n",
    "    total_train = 0\n",
    "\n",
    "    print(f\"Epoch [{epoch+1}/{EPOCHS}]\")\n",
    "    train_bar = tqdm(train_loader, desc=\"Training\", leave=False)\n",
    "\n",
    "    for images, labels in train_bar:\n",
    "        images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        total_train += labels.size(0)\n",
    "        correct_train += (predicted == labels).sum().item()\n",
    "\n",
    "        train_bar.set_postfix(loss=f\"{running_loss/len(train_loader):.4f}\", acc=f\"{100 * correct_train / total_train:.2f}%\")\n",
    "\n",
    "    train_accuracy = 100 * correct_train / total_train\n",
    "    print(f\"Train Loss: {running_loss/len(train_loader):.4f}, Train Accuracy: {train_accuracy:.2f}%\")\n",
    "\n",
    "    model.eval()\n",
    "    correct_test = 0\n",
    "    total_test = 0\n",
    "    test_loss = 0.0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        test_bar = tqdm(test_loader, desc=\"Testing\", leave=False)\n",
    "        for images, labels in test_bar:\n",
    "            images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total_test += labels.size(0)\n",
    "            correct_test += (predicted == labels).sum().item()\n",
    "\n",
    "    test_accuracy = 100 * correct_test / total_test\n",
    "    print(f\"Test Loss: {test_loss/len(test_loader):.4f}, Test Accuracy: {test_accuracy:.2f}%\")\n",
    "\n",
    "    if test_accuracy > best_accuracy:\n",
    "        best_accuracy = test_accuracy\n",
    "        torch.save(model.state_dict(), SAVE_PATH)\n",
    "        print(f\"Best model saved with test accuracy: {best_accuracy:.2f}%\")\n",
    "\n",
    "print(\"Обучение завершено.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_channels != out_channels:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),\n",
    "                nn.BatchNorm2d(out_channels)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = self.relu(out)\n",
    "        return out\n",
    "\n",
    "class EnhancedCNN(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(EnhancedCNN, self).__init__()\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "            nn.LeakyReLU(0.1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.25),\n",
    "\n",
    "            ResidualBlock(32, 64),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.25),\n",
    "\n",
    "            ResidualBlock(64, 128),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Dropout2d(0.25),\n",
    "\n",
    "            nn.AdaptiveAvgPool2d((1, 1))\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Linear(128, 256),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(256, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc_layers(x)\n",
    "        return x\n",
    "\n",
    "model = EnhancedCNN(len(CLASSES)).to(DEVICE)\n",
    "criterion = nn.CrossEntropyLoss()  # Или FocalLoss(), если необходимо\n",
    "optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch [1/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                   \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 2.9490, Train Accuracy: 9.62%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 2.2683, Test Accuracy: 31.25%\n",
      "Best model saved with test accuracy: 31.25%\n",
      "Улучшение найдено: 31.2500. Сохранение модели...\n",
      "\n",
      "Epoch [2/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 1.6123, Train Accuracy: 48.72%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.4703, Test Accuracy: 89.68%\n",
      "Best model saved with test accuracy: 89.68%\n",
      "Улучшение найдено: 89.6780. Сохранение модели...\n",
      "\n",
      "Epoch [3/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.6699, Train Accuracy: 79.42%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.2471, Test Accuracy: 93.47%\n",
      "Best model saved with test accuracy: 93.47%\n",
      "Улучшение найдено: 93.4659. Сохранение модели...\n",
      "\n",
      "Epoch [4/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.4096, Train Accuracy: 87.37%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.1863, Test Accuracy: 95.03%\n",
      "Best model saved with test accuracy: 95.03%\n",
      "Улучшение найдено: 95.0284. Сохранение модели...\n",
      "\n",
      "Epoch [5/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.2973, Train Accuracy: 91.12%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.1082, Test Accuracy: 96.83%\n",
      "Best model saved with test accuracy: 96.83%\n",
      "Улучшение найдено: 96.8277. Сохранение модели...\n",
      "\n",
      "Epoch [6/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.2404, Train Accuracy: 92.58%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.1003, Test Accuracy: 97.44%\n",
      "Best model saved with test accuracy: 97.44%\n",
      "Улучшение найдено: 97.4432. Сохранение модели...\n",
      "\n",
      "Epoch [7/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.1923, Train Accuracy: 94.24%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.0840, Test Accuracy: 97.63%\n",
      "Best model saved with test accuracy: 97.63%\n",
      "Улучшение найдено: 97.6326. Сохранение модели...\n",
      "\n",
      "Epoch [8/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.1699, Train Accuracy: 94.70%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.0862, Test Accuracy: 97.59%\n",
      "Ранняя остановка через 9 эпох(и).\n",
      "\n",
      "Epoch [9/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.1481, Train Accuracy: 95.53%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.0703, Test Accuracy: 97.54%\n",
      "Ранняя остановка через 8 эпох(и).\n",
      "\n",
      "Epoch [10/10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                    \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Loss: 0.1311, Train Accuracy: 95.96%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                        "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 0.0719, Test Accuracy: 97.82%\n",
      "Best model saved with test accuracy: 97.82%\n",
      "Улучшение найдено: 97.8220. Сохранение модели...\n",
      "Обучение завершено.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Предполагается, что следующие переменные уже определены:\n",
    "# model, train_loader, test_loader, DEVICE, EPOCHS, SAVE_PATH, CLASSES\n",
    "\n",
    "# Определение функции ранней остановки\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=10, verbose=False, delta=0):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            patience (int): Количество эпох без улучшения перед остановкой.\n",
    "            verbose (bool): Выводить сообщения о сохранении модели.\n",
    "            delta (float): Минимальное улучшение для учёта.\n",
    "        \"\"\"\n",
    "        self.patience = patience\n",
    "        self.verbose = verbose\n",
    "        self.delta = delta\n",
    "        self.best_score = None\n",
    "        self.early_stop = False\n",
    "        self.best_accuracy = 0.0\n",
    "\n",
    "    def __call__(self, score, model):\n",
    "        if self.best_score is None:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(score, model)\n",
    "        elif score < self.best_score + self.delta:\n",
    "            self.patience -= 1\n",
    "            if self.verbose:\n",
    "                print(f\"Ранняя остановка через {self.patience} эпох(и).\")\n",
    "            if self.patience <= 0:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_score = score\n",
    "            self.save_checkpoint(score, model)\n",
    "            self.patience = 10  # Сброс patience\n",
    "\n",
    "    def save_checkpoint(self, score, model):\n",
    "        if self.verbose:\n",
    "            print(f\"Улучшение найдено: {score:.4f}. Сохранение модели...\")\n",
    "        torch.save(model.state_dict(), SAVE_PATH)\n",
    "        self.best_accuracy = score\n",
    "\n",
    "# Инициализация функции потерь, оптимизатора и планировщика\n",
    "criterion = nn.CrossEntropyLoss()  # Или используйте FocalLoss(), если необходимо\n",
    "optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
    "early_stopping = EarlyStopping(patience=10, verbose=True)\n",
    "\n",
    "best_accuracy = 0.0\n",
    "\n",
    "for epoch in range(1, EPOCHS + 1):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    correct_train = 0\n",
    "    total_train = 0\n",
    "\n",
    "    print(f\"\\nEpoch [{epoch}/{EPOCHS}]\")\n",
    "    train_bar = tqdm(train_loader, desc=\"Training\", leave=False)\n",
    "\n",
    "    for images, labels in train_bar:\n",
    "        images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        total_train += labels.size(0)\n",
    "        correct_train += (predicted == labels).sum().item()\n",
    "\n",
    "        current_loss = running_loss / (train_bar.n + 1)\n",
    "        current_acc = 100 * correct_train / total_train\n",
    "        train_bar.set_postfix(loss=f\"{current_loss:.4f}\", acc=f\"{current_acc:.2f}%\")\n",
    "\n",
    "    train_loss = running_loss / len(train_loader)\n",
    "    train_accuracy = 100 * correct_train / total_train\n",
    "    print(f\"Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.2f}%\")\n",
    "\n",
    "    # Валидация\n",
    "    model.eval()\n",
    "    correct_test = 0\n",
    "    total_test = 0\n",
    "    test_loss = 0.0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        test_bar = tqdm(test_loader, desc=\"Testing\", leave=False)\n",
    "        for images, labels in test_bar:\n",
    "            images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            test_loss += loss.item()\n",
    "\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total_test += labels.size(0)\n",
    "            correct_test += (predicted == labels).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader)\n",
    "    test_accuracy = 100 * correct_test / total_test\n",
    "    print(f\"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%\")\n",
    "\n",
    "    # Обновление планировщика\n",
    "    scheduler.step()\n",
    "\n",
    "    # Проверка и обновление лучших результатов\n",
    "    if test_accuracy > best_accuracy:\n",
    "        best_accuracy = test_accuracy\n",
    "        torch.save(model.state_dict(), SAVE_PATH)\n",
    "        print(f\"Best model saved with test accuracy: {best_accuracy:.2f}%\")\n",
    "\n",
    "    # Применение ранней остановки\n",
    "    early_stopping(test_accuracy, model)\n",
    "    if early_stopping.early_stop:\n",
    "        print(\"Ранняя остановка активирована. Обучение завершено.\")\n",
    "        break\n",
    "\n",
    "print(\"Обучение завершено.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Error Statistics: defaultdict(<class 'int'>, {6: 19, 12: 7, 1: 8, 10: 8, 7: 4, 20: 8, 3: 6, 2: 8, 9: 4, 19: 9, 11: 8, 16: 9, 17: 8, 0: 5, 13: 4, 18: 8, 21: 2, 8: 1, 5: 5, 4: 7, 14: 1})\n",
      "Train Error Percentages: {6: 2.5299600532623168, 12: 0.8962868117797695, 1: 1.0062893081761006, 10: 1.0296010296010296, 7: 0.5161290322580645, 20: 1.0430247718383312, 3: 0.7905138339920948, 2: 1.06951871657754, 9: 0.5221932114882507, 19: 1.1583011583011582, 11: 1.0309278350515463, 16: 1.1450381679389312, 17: 1.0610079575596816, 0: 0.628140703517588, 13: 0.5263157894736842, 18: 1.0526315789473684, 21: 0.26246719160104987, 8: 0.13315579227696406, 5: 0.6459948320413437, 4: 0.9283819628647214, 14: 0.12919896640826875}\n",
      "Test Error Statistics: defaultdict(<class 'int'>, {10: 2, 2: 3, 5: 3, 16: 4, 17: 4, 20: 8, 19: 3, 14: 1, 6: 5, 11: 3, 13: 4, 0: 3, 1: 5, 15: 2, 21: 3, 4: 3, 9: 4, 18: 2, 7: 3, 12: 2})\n",
      "Test Error Percentages: {10: 1.092896174863388, 2: 1.4150943396226416, 5: 1.6129032258064515, 16: 2.2988505747126435, 17: 1.9417475728155338, 20: 4.145077720207254, 19: 1.639344262295082, 14: 0.5376344086021506, 6: 2.3923444976076556, 11: 1.6304347826086956, 13: 2.0, 0: 1.8292682926829267, 1: 3.0303030303030303, 15: 0.966183574879227, 21: 1.5151515151515151, 4: 1.4563106796116505, 9: 2.0618556701030926, 18: 1.0, 7: 1.6216216216216217, 12: 1.1173184357541899}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from collections import defaultdict\n",
    "\n",
    "def calculate_error_statistics(model, loader, device):\n",
    "    \"\"\"\n",
    "    Вычисляет общую статистику ошибок для датасета (train_loader/test_loader).\n",
    "\n",
    "    :param model: обученная модель\n",
    "    :param loader: DataLoader для вычисления статистики\n",
    "    :param device: устройство ('cuda' или 'cpu')\n",
    "    :return: словарь с количеством ошибок для каждого класса\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    error_stats = defaultdict(int)  # Хранение ошибок для каждого класса\n",
    "    total_samples = defaultdict(int)  # Общее количество образцов для каждого класса\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, targets in loader:\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            outputs = model(inputs)\n",
    "            predictions = torch.argmax(outputs, dim=1)\n",
    "\n",
    "            for pred, target in zip(predictions, targets):\n",
    "                total_samples[target.item()] += 1\n",
    "                if pred.item() != target.item():\n",
    "                    error_stats[target.item()] += 1\n",
    "\n",
    "    error_percentages = {\n",
    "        cls: (errors / total_samples[cls] * 100) if total_samples[cls] > 0 else 0\n",
    "        for cls, errors in error_stats.items()\n",
    "    }\n",
    "\n",
    "    return error_stats, error_percentages\n",
    "\n",
    "train_errors, train_error_percentages = calculate_error_statistics(model, train_loader, DEVICE)\n",
    "test_errors, test_error_percentages = calculate_error_statistics(model, test_loader, DEVICE)\n",
    "\n",
    "print(\"Train Error Statistics:\", train_errors)\n",
    "print(\"Train Error Percentages:\", train_error_percentages)\n",
    "print(\"Test Error Statistics:\", test_errors)\n",
    "print(\"Test Error Percentages:\", test_error_percentages)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47668 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# Функция для упорядочивания точек в определенном порядке: верхний левый, верхний правый, нижний правый, нижний левый\n",
    "def order_points(pts):\n",
    "    \"\"\"\n",
    "    Упорядочивает точки в порядке:\n",
    "    верхний левый, верхний правый, нижний правый, нижний левый\n",
    "    \"\"\"\n",
    "    rect = np.zeros((4, 2), dtype=\"float32\")\n",
    "\n",
    "    # Сумма координат для нахождения верхнего левого и нижнего правого углов\n",
    "    s = pts.sum(axis=1)\n",
    "    rect[0] = pts[np.argmin(s)]  # верхний левый\n",
    "    rect[2] = pts[np.argmax(s)]  # нижний правый\n",
    "\n",
    "    # Разница координат для нахождения верхнего правого и нижнего левого углов\n",
    "    diff = np.diff(pts, axis=1)\n",
    "    rect[1] = pts[np.argmin(diff)]  # верхний правый\n",
    "    rect[3] = pts[np.argmax(diff)]  # нижний левый\n",
    "\n",
    "    return rect\n",
    "\n",
    "# Функция для получения перспективного преобразования\n",
    "def get_transform(image, pts):\n",
    "    \"\"\"\n",
    "    Применяет перспективное преобразование к области, определенной точками pts в изображении.\n",
    "    Возвращает обрезанное и преобразованное изображение номерного знака.\n",
    "    \"\"\"\n",
    "    rect = order_points(pts)\n",
    "    (tl, tr, br, bl) = rect\n",
    "\n",
    "    # Вычисляем ширину нового изображения\n",
    "    widthA = np.linalg.norm(br - bl)\n",
    "    widthB = np.linalg.norm(tr - tl)\n",
    "    maxWidth = max(int(widthA), int(widthB))\n",
    "\n",
    "    # Вычисляем высоту нового изображения\n",
    "    heightA = np.linalg.norm(tr - br)\n",
    "    heightB = np.linalg.norm(tl - bl)\n",
    "    maxHeight = max(int(heightA), int(heightB))\n",
    "\n",
    "    # Определяем точки назначения для преобразования\n",
    "    dst = np.array([\n",
    "        [0, 0],\n",
    "        [maxWidth - 1, 0],\n",
    "        [maxWidth - 1, maxHeight - 1],\n",
    "        [0, maxHeight - 1]\n",
    "    ], dtype=\"float32\")\n",
    "\n",
    "    # Вычисляем матрицу перспективного преобразования\n",
    "    M = cv2.getPerspectiveTransform(rect, dst)\n",
    "\n",
    "    # Применяем перспективное преобразование\n",
    "    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))\n",
    "\n",
    "    return warped\n",
    "\n",
    "# Функция для получения преобразованного номерного знака\n",
    "def get_transformed_plate(image, bbox):\n",
    "    \"\"\"\n",
    "    Выполняет перспективное преобразование для извлечения номерного знака из изображения.\n",
    "    \"\"\"\n",
    "    plate = get_transform(image, bbox)\n",
    "    return plate\n",
    "\n",
    "# Список путей к изображениям\n",
    "images = ['img/1.jpg', 'img/2.jpg', 'img/3.jpg']\n",
    "\n",
    "processed_plates = []\n",
    "\n",
    "for img_path in images:\n",
    "    image = cv2.imread(img_path)\n",
    "    \n",
    "    if image is None:\n",
    "        print(f\"Не удалось загрузить изображение по пути: {img_path}\")\n",
    "        processed_plates.append(None)\n",
    "        continue\n",
    "\n",
    "    # Преобразуем изображение в оттенки серого\n",
    "    img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    \n",
    "    # Применяем пороговое значение для выделения контуров\n",
    "    ret, thresh = cv2.threshold(img_gray, 100, 200, cv2.THRESH_TOZERO_INV)\n",
    "    \n",
    "    # Находим контуры на изображении\n",
    "    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    plate_found = False  # Флаг для отслеживания обнаружения номерного знака\n",
    "\n",
    "    for cnt in contours:\n",
    "        x, y, w, h = cv2.boundingRect(cnt)\n",
    "        area = w * h\n",
    "        aspectRatio = float(w) / h\n",
    "\n",
    "        # Фильтруем контуры по соотношению сторон и площади\n",
    "        if aspectRatio >= 3 and area > 600:\n",
    "            approx = cv2.approxPolyDP(cnt, 0.05 * cv2.arcLength(cnt, True), True)\n",
    "            \n",
    "            # Проверяем, имеет ли контур 4 угла и находится ли он достаточно далеко от края\n",
    "            if len(approx) <= 4 and x > 15:\n",
    "                # Получаем минимальный ограничивающий прямоугольник с поворотом\n",
    "                rect = cv2.minAreaRect(cnt)\n",
    "                box = cv2.boxPoints(rect)\n",
    "                box = np.intp(box)\n",
    "\n",
    "                # Применяем перспективное преобразование для извлечения номерного знака\n",
    "                plate = get_transformed_plate(image, box)\n",
    "\n",
    "                processed_plates.append(plate)\n",
    "                plate_found = True\n",
    "                break  # Выходим из цикла после нахождения первого подходящего контура\n",
    "\n",
    "    if not plate_found:\n",
    "        print(f\"Номерной знак не найден на изображении: {img_path}\")\n",
    "        processed_plates.append(None)\n",
    "\n",
    "# Отображаем обработанные номерные знаки с помощью Matplotlib\n",
    "num_plates = len(processed_plates)\n",
    "\n",
    "# Определяем количество столбцов и строк для отображения\n",
    "cols = min(num_plates, 3)  # Ограничиваем до 3 столбцов\n",
    "rows = (num_plates + cols - 1) // cols  # Вычисляем необходимое количество строк\n",
    "\n",
    "fig, axs = plt.subplots(rows, cols, figsize=(5 * cols, 5 * rows))\n",
    "axs = axs.flatten() if num_plates > 1 else [axs]\n",
    "\n",
    "for i, rotated_plate in enumerate(processed_plates):\n",
    "    ax = axs[i]\n",
    "    if rotated_plate is not None:\n",
    "        # Проверяем, является ли изображение цветным\n",
    "        if len(rotated_plate.shape) == 3:\n",
    "            rotated_plate_rgb = cv2.cvtColor(rotated_plate, cv2.COLOR_BGR2RGB)\n",
    "            ax.imshow(rotated_plate_rgb)\n",
    "        else:\n",
    "            ax.imshow(rotated_plate, cmap='gray')\n",
    "    else:\n",
    "        # Отображаем сообщение, если номерной знак не найден\n",
    "        ax.text(0.5, 0.5, 'Номер не найден', horizontalalignment='center', verticalalignment='center', fontsize=12, color='red')\n",
    "    \n",
    "    ax.set_title(f'Номер {i+1}')\n",
    "    ax.axis('off')\n",
    "\n",
    "# Скрываем неиспользуемые подграфики, если они есть\n",
    "for j in range(i + 1, len(axs)):\n",
    "    fig.delaxes(axs[j])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57796 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def order_points(pts):\n",
    "    \"\"\"\n",
    "    Упорядочивает точки в порядке:\n",
    "    верхний левый, верхний правый, нижний правый, нижний левый\n",
    "    \"\"\"\n",
    "    rect = np.zeros((4, 2), dtype=\"float32\")\n",
    "\n",
    "    s = pts.sum(axis=1)\n",
    "    rect[0] = pts[np.argmin(s)] \n",
    "    rect[2] = pts[np.argmax(s)]\n",
    "\n",
    "    diff = np.diff(pts, axis=1)\n",
    "    rect[1] = pts[np.argmin(diff)] \n",
    "    rect[3] = pts[np.argmax(diff)] \n",
    "\n",
    "    return rect\n",
    "\n",
    "def get_transform(image, pts, padding=0.1):\n",
    "    \"\"\"\n",
    "    Применяет перспективное преобразование к области, определенной точками pts в изображении.\n",
    "    Уменьшает область захвата на заданный процент (padding).\n",
    "    Возвращает обрезанное и преобразованное изображение номерного знака.\n",
    "    \n",
    "    :param image: Исходное изображение\n",
    "    :param pts: Точки контура номерного знака (4 точки)\n",
    "    :param padding: Процент уменьшения области захвата (например, 0.1 для 10%)\n",
    "    \"\"\"\n",
    "    rect = order_points(pts)\n",
    "    (tl, tr, br, bl) = rect\n",
    "\n",
    "    widthA = np.linalg.norm(br - bl)\n",
    "    widthB = np.linalg.norm(tr - tl)\n",
    "    maxWidth = max(int(widthA), int(widthB))\n",
    "\n",
    "    heightA = np.linalg.norm(tr - br)\n",
    "    heightB = np.linalg.norm(tl - bl)\n",
    "    maxHeight = max(int(heightA), int(heightB))\n",
    "\n",
    "    center = np.mean(rect, axis=0)\n",
    "\n",
    "    vectors = rect - center\n",
    "\n",
    "    rect_shrinked = rect - vectors * padding\n",
    "\n",
    "    rect_shrinked[:, 0] = np.clip(rect_shrinked[:, 0], 0, image.shape[1] - 1)\n",
    "    rect_shrinked[:, 1] = np.clip(rect_shrinked[:, 1], 0, image.shape[0] - 1)\n",
    "\n",
    "    dst = np.array([\n",
    "        [0, 0],\n",
    "        [maxWidth - 1, 0],\n",
    "        [maxWidth - 1, maxHeight - 1],\n",
    "        [0, maxHeight - 1]\n",
    "    ], dtype=\"float32\")\n",
    "\n",
    "    M = cv2.getPerspectiveTransform(rect_shrinked, dst)\n",
    "\n",
    "    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))\n",
    "\n",
    "    return warped\n",
    "\n",
    "def get_transformed_plate(image, bbox, padding=0.1):\n",
    "    \"\"\"\n",
    "    Выполняет перспективное преобразование для извлечения номерного знака из изображения.\n",
    "    Уменьшает область захвата на заданный процент (padding).\n",
    "    \n",
    "    :param image: Исходное изображение\n",
    "    :param bbox: Координаты ограничивающего прямоугольника (4 точки)\n",
    "    :param padding: Процент уменьшения области захвата (например, 0.1 для 10%)\n",
    "    :return: Преобразованное изображение номерного знака\n",
    "    \"\"\"\n",
    "    plate = get_transform(image, bbox, padding)\n",
    "    return plate\n",
    "\n",
    "images = ['img/1.jpg', 'img/2.jpg', 'img/3.jpg', 'img/ru4018185.jpg']\n",
    "\n",
    "processed_plates = []\n",
    "\n",
    "for img_path in images:\n",
    "    image = cv2.imread(img_path)\n",
    "    \n",
    "    if image is None:\n",
    "        print(f\"Не удалось загрузить изображение по пути: {img_path}\")\n",
    "        processed_plates.append(None)\n",
    "        continue\n",
    "\n",
    "    img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    \n",
    "    ret, thresh = cv2.threshold(img_gray, 100, 200, cv2.THRESH_TOZERO_INV)\n",
    "    \n",
    "    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    plate_found = False \n",
    "\n",
    "    for cnt in contours:\n",
    "        x, y, w, h = cv2.boundingRect(cnt)\n",
    "        area = w * h\n",
    "        aspectRatio = float(w) / h\n",
    "\n",
    "        if aspectRatio >= 3 and area > 600:\n",
    "            approx = cv2.approxPolyDP(cnt, 0.05 * cv2.arcLength(cnt, True), True)\n",
    "            \n",
    "            if len(approx) <= 4 and x > 15:\n",
    "                rect = cv2.minAreaRect(cnt)\n",
    "                box = cv2.boxPoints(rect)\n",
    "                box = np.intp(box)\n",
    "\n",
    "                \n",
    "                plate = get_transformed_plate(image, box, padding=0.1) \n",
    "\n",
    "                processed_plates.append(plate)\n",
    "                plate_found = True\n",
    "                break \n",
    "\n",
    "    if not plate_found:\n",
    "        print(f\"Номерной знак не найден на изображении: {img_path}\")\n",
    "        processed_plates.append(None)\n",
    "\n",
    "num_plates = len(processed_plates)\n",
    "\n",
    "cols = min(num_plates, 3)\n",
    "rows = (num_plates + cols - 1) // cols \n",
    "\n",
    "fig, axs = plt.subplots(rows, cols, figsize=(5 * cols, 5 * rows))\n",
    "if num_plates == 1:\n",
    "    axs = [axs] \n",
    "else:\n",
    "    axs = axs.flatten()\n",
    "\n",
    "for i, rotated_plate in enumerate(processed_plates):\n",
    "    ax = axs[i]\n",
    "    if rotated_plate is not None:\n",
    "        if len(rotated_plate.shape) == 3:\n",
    "            rotated_plate_rgb = cv2.cvtColor(rotated_plate, cv2.COLOR_BGR2RGB)\n",
    "            ax.imshow(rotated_plate_rgb)\n",
    "        else:\n",
    "            ax.imshow(rotated_plate, cmap='gray')\n",
    "    else:\n",
    "        ax.text(0.5, 0.5, 'Номер не найден', horizontalalignment='center', verticalalignment='center', fontsize=12, color='red')\n",
    "    \n",
    "    ax.set_title(f'Номер {i+1}')\n",
    "    ax.axis('off')\n",
    "\n",
    "for j in range(i + 1, len(axs)):\n",
    "    fig.delaxes(axs[j])\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1076228 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 975456 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 930660 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def order_points(pts):\n",
    "    rect = np.zeros((4, 2), dtype=\"float32\")\n",
    "    s = pts.sum(axis=1)\n",
    "    rect[0] = pts[np.argmin(s)]\n",
    "    rect[2] = pts[np.argmax(s)]\n",
    "    diff = np.diff(pts, axis=1)\n",
    "    rect[1] = pts[np.argmin(diff)]\n",
    "    rect[3] = pts[np.argmax(diff)]\n",
    "    return rect\n",
    "\n",
    "def get_transform(image, pts, padding=0.1):\n",
    "    rect = order_points(pts)\n",
    "    (tl, tr, br, bl) = rect\n",
    "    widthA = np.linalg.norm(br - bl)\n",
    "    widthB = np.linalg.norm(tr - tl)\n",
    "    maxWidth = max(int(widthA), int(widthB))\n",
    "    heightA = np.linalg.norm(tr - br)\n",
    "    heightB = np.linalg.norm(tl - bl)\n",
    "    maxHeight = max(int(heightA), int(heightB))\n",
    "    center = np.mean(rect, axis=0)\n",
    "    vectors = rect - center\n",
    "    rect_shrinked = rect - vectors * padding\n",
    "    rect_shrinked[:, 0] = np.clip(rect_shrinked[:, 0], 0, image.shape[1] - 1)\n",
    "    rect_shrinked[:, 1] = np.clip(rect_shrinked[:, 1], 0, image.shape[0] - 1)\n",
    "    dst = np.array([\n",
    "        [0, 0],\n",
    "        [maxWidth - 1, 0],\n",
    "        [maxWidth - 1, maxHeight - 1],\n",
    "        [0, maxHeight - 1]\n",
    "    ], dtype=\"float32\")\n",
    "    M = cv2.getPerspectiveTransform(rect_shrinked, dst)\n",
    "    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))\n",
    "    return warped\n",
    "\n",
    "def get_transformed_plate(image, bbox, padding=0.1):\n",
    "    plate = get_transform(image, bbox, padding)\n",
    "    return plate\n",
    "\n",
    "def process_and_display_plate(image_path, padding=0.1):\n",
    "    image = cv2.imread(image_path)\n",
    "    if image is None:\n",
    "        print(f\"Failed to load image at path: {image_path}\")\n",
    "        return\n",
    "    output_image = image.copy()\n",
    "    img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    ret, thresh = cv2.threshold(img_gray, 100, 200, cv2.THRESH_TOZERO_INV)\n",
    "    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    plate_found = False\n",
    "    plate_image = None\n",
    "    plate_box = None\n",
    "    for cnt in contours:\n",
    "        x, y, w, h = cv2.boundingRect(cnt)\n",
    "        area = w * h\n",
    "        aspectRatio = float(w) / h\n",
    "        if aspectRatio >= 3 and area > 600:\n",
    "            approx = cv2.approxPolyDP(cnt, 0.05 * cv2.arcLength(cnt, True), True)\n",
    "            if len(approx) <= 4 and x > 15:\n",
    "                rect = cv2.minAreaRect(cnt)\n",
    "                box = cv2.boxPoints(rect)\n",
    "                box = np.intp(box)\n",
    "                plate_image = get_transformed_plate(image, box, padding)\n",
    "                plate_found = True\n",
    "                plate_box = box\n",
    "                break\n",
    "    if plate_found:\n",
    "        cv2.polylines(output_image, [plate_box], isClosed=True, color=(0, 255, 0), thickness=2)\n",
    "    height, width = image.shape[:2]\n",
    "    black_space_width = width // 2\n",
    "    black_space = np.zeros((height, black_space_width, 3), dtype=np.uint8)\n",
    "    if plate_found and plate_image is not None:\n",
    "        plate_height, plate_width = plate_image.shape[:2]\n",
    "        scale_factor = min(black_space_width / (plate_width + 200), height / plate_height)\n",
    "        new_plate_width = int(plate_width * scale_factor)\n",
    "        new_plate_height = int(plate_height * scale_factor)\n",
    "        resized_plate = cv2.resize(plate_image, (new_plate_width, new_plate_height), interpolation=cv2.INTER_AREA)\n",
    "        y_offset = (height - new_plate_height) // 2\n",
    "        black_space[y_offset:y_offset + new_plate_height, 0:new_plate_width] = resized_plate\n",
    "        placeholder_text = \"Text plate:\"\n",
    "        decoded_text = \"__________\"\n",
    "        font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "        font_scale = 1\n",
    "        font_color = (255, 255, 255)\n",
    "        thickness = 2\n",
    "        line_type = cv2.LINE_AA\n",
    "        cv2.putText(black_space, placeholder_text, (new_plate_width + 20, y_offset + new_plate_height // 2), \n",
    "                    font, font_scale, font_color, thickness, line_type)\n",
    "        cv2.putText(black_space, decoded_text, (new_plate_width + 20, y_offset + new_plate_height // 2 + 50), \n",
    "                    font, font_scale, font_color, thickness, line_type)\n",
    "    else:\n",
    "        placeholder_text = \"Not found.\"\n",
    "        font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "        font_scale = 1\n",
    "        font_color = (0, 0, 255)\n",
    "        thickness = 2\n",
    "        line_type = cv2.LINE_AA\n",
    "        text_size, _ = cv2.getTextSize(placeholder_text, font, font_scale, thickness)\n",
    "        text_x = (black_space_width - text_size[0]) // 2\n",
    "        text_y = height // 2\n",
    "        cv2.putText(black_space, placeholder_text, (text_x, text_y), \n",
    "                    font, font_scale, font_color, thickness, line_type)\n",
    "    combined_image = np.hstack((output_image, black_space))\n",
    "    combined_image_rgb = cv2.cvtColor(combined_image, cv2.COLOR_BGR2RGB)\n",
    "    plt.figure(figsize=(15, 10))\n",
    "    plt.imshow(combined_image_rgb)\n",
    "    plt.axis('off')\n",
    "    plt.title(f'Processed Image: {image_path}')\n",
    "    plt.show()\n",
    "\n",
    "images = ['img/1.jpg', 'img/2.jpg', 'img/3.jpg']\n",
    "\n",
    "for img_path in images:\n",
    "    process_and_display_plate(img_path, padding=0.1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 98304 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 108168 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 101944 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# Функция удаления мелких объектов\n",
    "def bwareaopen(imgBW, areaPixels):\n",
    "    imgBWcopy = imgBW.copy()\n",
    "    contours, _ = cv2.findContours(imgBWcopy, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    for contour in contours:\n",
    "        if cv2.contourArea(contour) < areaPixels:\n",
    "            cv2.drawContours(imgBWcopy, [contour], -1, 0, -1)\n",
    "    return imgBWcopy\n",
    "\n",
    "# Функция расширения объектов\n",
    "def expand_characters(img, kernel_size=(2, 2), iterations=1):\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)\n",
    "    return cv2.dilate(img, kernel, iterations=iterations)\n",
    "\n",
    "def selective_dilation(image, distance_threshold=10, iterations=1):\n",
    "    \"\"\"\n",
    "    Выполняет \"наращивание\" (дилатацию) символов, чтобы соединить близкие контуры.\n",
    "    \n",
    "    :param image: Бинарное изображение.\n",
    "    :param distance_threshold: Максимальное расстояние между контурами для их соединения.\n",
    "    :param iterations: Количество итераций для расширения.\n",
    "    :return: Обработанное изображение.\n",
    "    \"\"\"\n",
    "    # Копируем изображение\n",
    "    image_copy = image.copy()\n",
    "\n",
    "    # Находим контуры\n",
    "    contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    # Создаём маску для временного соединения контуров\n",
    "    temp_mask = np.zeros_like(image)\n",
    "\n",
    "    for i, cnt1 in enumerate(contours):\n",
    "        for j, cnt2 in enumerate(contours):\n",
    "            if i != j:  # Сравниваем контуры только с другими\n",
    "                # Находим минимальное расстояние между точками двух контуров\n",
    "                min_distance = float(\"inf\")\n",
    "                for pt1 in cnt1:\n",
    "                    for pt2 in cnt2:\n",
    "                        dist = np.linalg.norm(np.array(pt1[0]) - np.array(pt2[0]))\n",
    "                        if dist < min_distance:\n",
    "                            min_distance = dist\n",
    "\n",
    "                # Если расстояние меньше порога, соединяем контуры\n",
    "                if min_distance < distance_threshold:\n",
    "                    # Рисуем оба контура на временной маске\n",
    "                    cv2.drawContours(temp_mask, [cnt1], -1, 255, thickness=-1)\n",
    "                    cv2.drawContours(temp_mask, [cnt2], -1, 255, thickness=-1)\n",
    "\n",
    "    # Выполняем дилатацию только на временной маске\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))\n",
    "    dilated_mask = cv2.dilate(temp_mask, kernel, iterations=iterations)\n",
    "\n",
    "    # Объединяем исходное изображение с обработанной маской\n",
    "    result = cv2.bitwise_or(image_copy, dilated_mask)\n",
    "\n",
    "    return result\n",
    "\n",
    "def imclearborder(imgBW, radius=1):\n",
    "    \"\"\"\n",
    "    Удаляет только те объекты, которые непосредственно касаются границ изображения,\n",
    "    если у них нет точек в пределах 30% от верхней или нижней границы.\n",
    "    \n",
    "    :param imgBW: Бинарное изображение (чёрно-белое, 0 и 255).\n",
    "    :param radius: Расстояние от границы, в пределах которого объект считается касающимся.\n",
    "    :return: Изображение с удалёнными объектами, касающимися границ.\n",
    "    \"\"\"\n",
    "    imgBWcopy = imgBW.copy()\n",
    "\n",
    "    # Находим контуры\n",
    "    contours, _ = cv2.findContours(imgBWcopy.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    imgRows, imgCols = imgBW.shape\n",
    "\n",
    "    # Границы для проверки: 30% сверху и снизу\n",
    "    top_limit = int(imgRows * 0.3)\n",
    "    bottom_limit = int(imgRows * 0.7)\n",
    "\n",
    "    # Список контуров, которые касаются границы\n",
    "    contourList = []\n",
    "\n",
    "    for idx, cnt in enumerate(contours):\n",
    "        remove_contour = True  # Предположим, что контур нужно удалить\n",
    "        for pt in cnt:\n",
    "            rowCnt = pt[0][1]  # y-координата\n",
    "            colCnt = pt[0][0]  # x-координата\n",
    "\n",
    "            # Проверка на касание границы\n",
    "            check1 = (rowCnt >= 0 and rowCnt < radius) or (rowCnt >= imgRows - radius and rowCnt < imgRows)\n",
    "            check2 = (colCnt >= 0 and colCnt < radius) or (colCnt >= imgCols - radius and colCnt < imgCols)\n",
    "\n",
    "            # Если точка находится в пределах 30% от верха или низа, не удаляем контур\n",
    "            if top_limit <= rowCnt <= bottom_limit:\n",
    "                remove_contour = False\n",
    "                break\n",
    "\n",
    "            if check1 or check2:\n",
    "                continue  # Проверяем следующую точку\n",
    "\n",
    "        if remove_contour:\n",
    "            contourList.append(idx)\n",
    "\n",
    "    # Удаляем найденные контуры\n",
    "    for idx in contourList:\n",
    "        cv2.drawContours(imgBWcopy, contours, idx, (0, 0, 0), -1)\n",
    "\n",
    "    return imgBWcopy\n",
    "\n",
    "\n",
    "\n",
    "# Основная функция обработки одной пластины\n",
    "def process_plate(plate):\n",
    "    # Преобразование в градации серого\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    # Увеличение размера для улучшения качества\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "\n",
    "    # Фильтрация и бинаризация\n",
    "    gray = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "    thresh = cv2.adaptiveThreshold(gray, 255,\n",
    "                                   cv2.ADAPTIVE_THRESH_GAUSSIAN_C,\n",
    "                                   cv2.THRESH_BINARY_INV, 41, 10)\n",
    "\n",
    "    # Удаление мелких областей\n",
    "    #expanded_image = expand_characters(thresh, kernel_size=(2, 2), iterations=1)\n",
    "    thresh = imclearborder(thresh, 5)\n",
    "    expanded_image = selective_dilation(thresh, distance_threshold=7, iterations=1)\n",
    "    processed_plate = bwareaopen(expanded_image, 20)\n",
    "\n",
    "    # Поиск контуров символов\n",
    "    contours, _ = cv2.findContours(processed_plate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    # Сортировка контуров по их расположению (слева направо)\n",
    "    contours = sorted(contours, key=lambda ctr: cv2.boundingRect(ctr)[0])\n",
    "\n",
    "    # Вырезание и нормализация каждого символа\n",
    "    characters = []\n",
    "    for ctr in contours:\n",
    "        x, y, w, h = cv2.boundingRect(ctr)\n",
    "        if w > 2 and h > 10:  # Убираем шумы и маленькие контуры\n",
    "            char = processed_plate[y:y+h, x:x+w]\n",
    "            char_resized = cv2.resize(char, (28, 28))  # Подгонка под размер (например, 28x28 для модели)\n",
    "            characters.append(char_resized)\n",
    "    return processed_plate, characters\n",
    "\n",
    "# Пример использования с массивом processed_plates\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "    # Обработка пластинки\n",
    "    processed_plate, characters = process_plate(plate)\n",
    "\n",
    "    # Вывод результата\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # Исходная пластина\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.title(\"Исходная картинка\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))  # Преобразуем для корректного отображения\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Обработанная пластина\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.title(\"Обработанная картинка\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.title(\"Выделенные символы\")\n",
    "    if characters:\n",
    "        symbols_grid = np.hstack(characters)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'Нет символов', ha='center', va='center')\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.suptitle(f\"Результат обработки пластинки {plate_index + 1}\")\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved Original Plate to img\\plate1_original.jpg\n",
      "Saved Processed Plate to img\\plate1_processed.jpg\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 107604 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved Original Plate to img\\plate2_original.jpg\n",
      "Saved Processed Plate to img\\plate2_processed.jpg\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 117308 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved Original Plate to img\\plate3_original.jpg\n",
      "Saved Processed Plate to img\\plate3_processed.jpg\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 107816 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image, ImageOps  # Added for padding\n",
    "\n",
    "# 1. Function to add padding around symbols\n",
    "def add_padding(img, padding=10):\n",
    "    \"\"\"\n",
    "    Adds padding around the image.\n",
    "\n",
    "    :param img: Numpy array with the symbol image.\n",
    "    :param padding: Number of pixels for padding.\n",
    "    :return: Image with padding.\n",
    "    \"\"\"\n",
    "    # Ensure the image is in uint8 format\n",
    "    if img.dtype != np.uint8:\n",
    "        img = img.astype(np.uint8)\n",
    "    \n",
    "    pil_img = Image.fromarray(img)\n",
    "    padded_img = ImageOps.expand(pil_img, border=padding, fill=255)\n",
    "    return np.array(padded_img)\n",
    "\n",
    "# Define mask positions based on GOST standards\n",
    "# These masks are defined as fractions of the image width and height\n",
    "# Adjust these values based on the actual license plate layout and resolution\n",
    "\n",
    "# Mask for plates with 2 region digits (total 8 symbols)\n",
    "mask_8 = [\n",
    "    (0.05, 0.25, 0.1, 0.5),  # Символ 1: Буква\n",
    "    (0.17, 0.25, 0.1, 0.5),  # Символ 2: Цифра\n",
    "    (0.29, 0.25, 0.1, 0.5),  # Символ 3: Цифра\n",
    "    (0.41, 0.25, 0.1, 0.5),  # Символ 4: Цифра\n",
    "    (0.53, 0.25, 0.1, 0.5),  # Символ 5: Буква\n",
    "    (0.65, 0.25, 0.1, 0.5),  # Символ 6: Буква\n",
    "    (0.80, 0.05, 0.1, 0.5), # Символ 7: Региональная цифра 1 (выше и меньше)\n",
    "    (0.90, 0.05, 0.1, 0.5),  # Символ 8: Региональная цифра 2 (выше и меньше)\n",
    "]\n",
    "\n",
    "\n",
    "# Mask for plates with 3 region digits (total 9 symbols)\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),  # Symbol 1: Letter\n",
    "    (0.17, 0.2, 0.1, 0.6),  # Symbol 2: Digit\n",
    "    (0.29, 0.2, 0.1, 0.6),  # Symbol 3: Digit\n",
    "    (0.41, 0.2, 0.1, 0.6),  # Symbol 4: Digit\n",
    "    (0.53, 0.2, 0.1, 0.6),  # Symbol 5: Letter\n",
    "    (0.65, 0.2, 0.1, 0.6),  # Symbol 6: Letter\n",
    "    (0.77, 0.2, 0.1, 0.6),  # Symbol 7: Region Digit 1\n",
    "    (0.89, 0.2, 0.1, 0.6),  # Symbol 8: Region Digit 2\n",
    "    (1.01, 0.2, 0.1, 0.6)   # Symbol 9: Region Digit 3\n",
    "]\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image, ImageOps\n",
    "\n",
    "def resize_with_padding(image, size=28, padding=4, bg_color=0):\n",
    "    \"\"\"\n",
    "    Изменяет размер изображения до заданного размера с сохранением соотношения сторон\n",
    "    и добавлением отступов.\n",
    "    \n",
    "    :param image: Градация серого изображения символа (numpy array).\n",
    "    :param size: Целевой размер (ширина и высота) итогового изображения.\n",
    "    :param padding: Количество пикселей для отступов.\n",
    "    :param bg_color: Цвет фона (0 для чёрного, 255 для белого).\n",
    "    :return: Изображение размером size x size пикселей.\n",
    "    \"\"\"\n",
    "    # Преобразуем numpy array в PIL Image\n",
    "    pil_img = Image.fromarray(image)\n",
    "    \n",
    "    # Определяем текущие размеры\n",
    "    w, h = pil_img.size\n",
    "    \n",
    "    # Рассчитываем масштабирование\n",
    "    scale = (size - 2 * padding) / max(w, h)\n",
    "    new_w = int(w * scale)\n",
    "    new_h = int(h * scale)\n",
    "    \n",
    "    # Изменяем размер с сохранением соотношения сторон\n",
    "    pil_resized = pil_img.resize((new_w, new_h), Image.LANCZOS)\n",
    "    \n",
    "    # Создаём новое изображение с фоном\n",
    "    new_img = Image.new(\"L\", (size, size), color=bg_color)\n",
    "    \n",
    "    # Вычисляем позицию для центрирования\n",
    "    paste_x = (size - new_w) // 2\n",
    "    paste_y = (size - new_h) // 2\n",
    "    \n",
    "    # Вставляем изменённое изображение на фон\n",
    "    new_img.paste(pil_resized, (paste_x, paste_y))\n",
    "    \n",
    "    # Преобразуем обратно в numpy array\n",
    "    return np.array(new_img)\n",
    "\n",
    "\n",
    "\n",
    "def extract_symbols(image, masks):\n",
    "    \"\"\"\n",
    "    Extracts symbols from the image based on the provided masks.\n",
    "\n",
    "    :param image: Grayscale image of the license plate.\n",
    "    :param masks: List of tuples containing mask definitions as (x_frac, y_frac, w_frac, h_frac).\n",
    "    :return: List of extracted and resized symbol images.\n",
    "    \"\"\"\n",
    "    img_h, img_w = image.shape\n",
    "    symbols = []\n",
    "    for bbox in masks:\n",
    "        x_frac, y_frac, w_frac, h_frac = bbox\n",
    "        x = int(x_frac * img_w)\n",
    "        y = int(y_frac * img_h)\n",
    "        w = int(w_frac * img_w)\n",
    "        h = int(h_frac * img_h)\n",
    "        \n",
    "        # Ensure the bounding box is within image boundaries\n",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        \n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            # If the bounding box is out of image bounds, skip\n",
    "            continue\n",
    "        # Add padding\n",
    "        # Resize to 28x28 pixels\n",
    "        symbol_resized = resize_with_padding(symbol, size=28, padding=2, bg_color=255)\n",
    "        symbols.append(symbol_resized)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    \"\"\"\n",
    "    Determines if a symbol is present based on the mean pixel intensity.\n",
    "\n",
    "    :param symbol: Grayscale image of the symbol.\n",
    "    :param threshold: Intensity threshold to consider a symbol as present.\n",
    "    :return: Boolean indicating presence of a symbol.\n",
    "    \"\"\"\n",
    "    return np.mean(symbol) < threshold  # Darker symbols have lower mean intensity\n",
    "\n",
    "# Main processing function\n",
    "def process_plate(plate):\n",
    "    \"\"\"\n",
    "    Processes a single license plate image to extract individual symbols based on masks.\n",
    "\n",
    "    :param plate: BGR image of the cropped license plate.\n",
    "    :return: Tuple containing the processed grayscale image and list of extracted symbols.\n",
    "    \"\"\"\n",
    "    # Convert to grayscale\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    # Resize for consistent processing (adjust fx and fy as needed)\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    img_h, img_w = resized_plate.shape\n",
    "\n",
    "    # Apply Gaussian blur to reduce noise\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "\n",
    "    # Extract symbols using both mask configurations\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\n",
    "\n",
    "    # Determine which mask to use based on the presence of the third region digit\n",
    "    # Check if the 9th symbol (Region Digit 3) is present\n",
    "    if len(symbols_9) == 9 and is_symbol_present(symbols_9[-1]):\n",
    "        characters = symbols_9\n",
    "    else:\n",
    "        characters = symbols_8\n",
    "\n",
    "    return blurred, characters\n",
    "\n",
    "# Example usage with processed_plates\n",
    "# Assume processed_plates is a list of BGR images containing cropped license plates\n",
    "# processed_plates = [cv2.imread('plate1.jpg'), cv2.imread('plate2.jpg'), ...]\n",
    "\n",
    "output_dir = \"img\"\n",
    "\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "    # Ensure the plate image is valid\n",
    "    if plate is None:\n",
    "        print(f\"Plate {plate_index + 1} could not be loaded.\")\n",
    "        continue\n",
    "\n",
    "    # Process the plate to extract symbols\n",
    "    processed_plate, characters = process_plate(plate)\n",
    "\n",
    "    # Сохранение Original Plate\n",
    "    original_path = os.path.join(output_dir, f\"plate{plate_index + 1}_original.jpg\")\n",
    "    cv2.imwrite(original_path, plate)\n",
    "    print(f\"Saved Original Plate to {original_path}\")\n",
    "\n",
    "    # Сохранение Processed Plate (Grayscale & Blurred)\n",
    "    processed_path = os.path.join(output_dir, f\"plate{plate_index + 1}_processed.jpg\")\n",
    "    cv2.imwrite(processed_path, processed_plate)\n",
    "    print(f\"Saved Processed Plate to {processed_path}\")\n",
    "\n",
    "    # Display the results\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # Original Plate\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.title(\"Original Plate\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Processed Plate (Grayscale & Blurred)\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.title(\"Processed Plate\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Extracted Symbols\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.title(\"Extracted Symbols\")\n",
    "    if characters:\n",
    "        # Arrange symbols in a horizontal grid\n",
    "        symbols_grid = np.hstack(characters)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.suptitle(f\"License Plate Processing Result {plate_index + 1}\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'SimpleCNN' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[1], line 135\u001b[0m\n\u001b[0;32m    126\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    127\u001b[0m \u001b[38;5;124;03m    Determines if a symbol is present based on the mean pixel intensity.\u001b[39;00m\n\u001b[0;32m    128\u001b[0m \n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    131\u001b[0m \u001b[38;5;124;03m    :return: Boolean indicating presence of a symbol.\u001b[39;00m\n\u001b[0;32m    132\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m    133\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m np\u001b[38;5;241m.\u001b[39mmean(symbol) \u001b[38;5;241m<\u001b[39m threshold  \u001b[38;5;66;03m# Darker symbols have lower mean intensity\u001b[39;00m\n\u001b[1;32m--> 135\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mSimpleCNN\u001b[49m(\u001b[38;5;28mlen\u001b[39m(CLASSES))\u001b[38;5;241m.\u001b[39mto(DEVICE)\n\u001b[0;32m    136\u001b[0m SAVE_PATH \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbest_model_5.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    137\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "\u001b[1;31mNameError\u001b[0m: name 'SimpleCNN' is not defined"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image, ImageOps  # Added for padding\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "\n",
    "# 1. Function to add padding around symbols\n",
    "def add_padding(img, padding=10):\n",
    "    \"\"\"\n",
    "    Adds padding around the image.\n",
    "\n",
    "    :param img: Numpy array with the symbol image.\n",
    "    :param padding: Number of pixels for padding.\n",
    "    :return: Image with padding.\n",
    "    \"\"\"\n",
    "    # Ensure the image is in uint8 format\n",
    "    if img.dtype != np.uint8:\n",
    "        img = img.astype(np.uint8)\n",
    "    \n",
    "    pil_img = Image.fromarray(img)\n",
    "    padded_img = ImageOps.expand(pil_img, border=padding, fill=255)\n",
    "    return np.array(padded_img)\n",
    "\n",
    "# Define mask positions based on GOST standards\n",
    "# These masks are defined as fractions of the image width and height\n",
    "# Adjust these values based on the actual license plate layout and resolution\n",
    "\n",
    "# Mask for plates with 2 region digits (total 8 symbols)\n",
    "mask_8 = [\n",
    "    (0.03, 0.05, 0.12, 0.9),  # Символ 1: Буква\n",
    "    (0.17, 0.05, 0.12, 0.9),  # Символ 2: Цифра\n",
    "    (0.29, 0.05, 0.12, 0.9),  # Символ 3: Цифра\n",
    "    (0.41, 0.05, 0.12, 0.9),  # Символ 4: Цифра\n",
    "    (0.53, 0.05, 0.12, 0.9),  # Символ 5: Буква\n",
    "    (0.65, 0.05, 0.12, 0.9),  # Символ 6: Буква\n",
    "    (0.80, 0.05, 0.11, 0.9),  # Символ 7: Региональная цифра 1 (выше и меньше)\n",
    "    (0.90, 0.05, 0.11, 0.9),  # Символ 8: Региональная цифра 2 (выше и меньше)\n",
    "]\n",
    "\n",
    "# Mask for plates with 3 region digits (total 9 symbols)\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),  # Symbol 1: Letter\n",
    "    (0.17, 0.2, 0.1, 0.6),  # Symbol 2: Digit\n",
    "    (0.29, 0.2, 0.1, 0.6),  # Symbol 3: Digit\n",
    "    (0.41, 0.2, 0.1, 0.6),  # Symbol 4: Digit\n",
    "    (0.53, 0.2, 0.1, 0.6),  # Symbol 5: Letter\n",
    "    (0.65, 0.2, 0.1, 0.6),  # Symbol 6: Letter\n",
    "    (0.77, 0.2, 0.1, 0.6),  # Symbol 7: Region Digit 1\n",
    "    (0.89, 0.2, 0.1, 0.6),  # Symbol 8: Region Digit 2\n",
    "    (1.01, 0.2, 0.1, 0.6)   # Symbol 9: Region Digit 3\n",
    "]\n",
    "\n",
    "def resize_with_padding(image, size=28, padding=4, bg_color=255):\n",
    "    \"\"\"\n",
    "    Изменяет размер изображения до заданного размера с сохранением соотношения сторон\n",
    "    и добавлением отступов.\n",
    "    \n",
    "    :param image: Градация серого изображения символа (numpy array).\n",
    "    :param size: Целевой размер (ширина и высота) итогового изображения.\n",
    "    :param padding: Количество пикселей для отступов.\n",
    "    :param bg_color: Цвет фона (0 для чёрного, 255 для белого).\n",
    "    :return: Изображение размером size x size пикселей.\n",
    "    \"\"\"\n",
    "    # Преобразуем numpy array в PIL Image\n",
    "    pil_img = Image.fromarray(image)\n",
    "    \n",
    "    # Определяем текущие размеры\n",
    "    w, h = pil_img.size\n",
    "    \n",
    "    # Рассчитываем масштабирование\n",
    "    scale = (size - 2 * padding) / max(w, h)\n",
    "    new_w = int(w * scale)\n",
    "    new_h = int(h * scale)\n",
    "    \n",
    "    # Изменяем размер с сохранением соотношения сторон\n",
    "    pil_resized = pil_img.resize((new_w, new_h), Image.LANCZOS)\n",
    "    \n",
    "    # Создаём новое изображение с фоном\n",
    "    new_img = Image.new(\"L\", (size, size), color=bg_color)\n",
    "    \n",
    "    # Вычисляем позицию для центрирования\n",
    "    paste_x = (size - new_w) // 2\n",
    "    paste_y = (size - new_h) // 2\n",
    "    \n",
    "    # Вставляем изменённое изображение на фон\n",
    "    new_img.paste(pil_resized, (paste_x, paste_y))\n",
    "    \n",
    "    # Преобразуем обратно в numpy array\n",
    "    return np.array(new_img)\n",
    "\n",
    "def extract_symbols(image, masks):\n",
    "    \"\"\"\n",
    "    Extracts symbols from the image based on the provided masks.\n",
    "\n",
    "    :param image: Grayscale image of the license plate.\n",
    "    :param masks: List of tuples containing mask definitions as (x_frac, y_frac, w_frac, h_frac).\n",
    "    :return: List of extracted and resized symbol images.\n",
    "    \"\"\"\n",
    "    img_h, img_w = image.shape\n",
    "    symbols = []\n",
    "    for bbox in masks:\n",
    "        x_frac, y_frac, w_frac, h_frac = bbox\n",
    "        x = int(x_frac * img_w)\n",
    "        y = int(y_frac * img_h)\n",
    "        w = int(w_frac * img_w)\n",
    "        h = int(h_frac * img_h)\n",
    "        \n",
    "        # Ensure the bounding box is within image boundaries\n",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        \n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            # If the bounding box is out of image bounds, skip\n",
    "            continue\n",
    "        # Add padding\n",
    "        # Resize to 28x28 pixels\n",
    "        symbol_resized = resize_with_padding(symbol, size=28, padding=2, bg_color=255)\n",
    "        symbols.append(symbol_resized)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    \"\"\"\n",
    "    Determines if a symbol is present based on the mean pixel intensity.\n",
    "\n",
    "    :param symbol: Grayscale image of the symbol.\n",
    "    :param threshold: Intensity threshold to consider a symbol as present.\n",
    "    :return: Boolean indicating presence of a symbol.\n",
    "    \"\"\"\n",
    "    return np.mean(symbol) < threshold  # Darker symbols have lower mean intensity\n",
    "\n",
    "model = SimpleCNN(len(CLASSES)).to(DEVICE)\n",
    "SAVE_PATH = \"best_model_5.pth\"\n",
    "try:\n",
    "    model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "    model.eval()\n",
    "    print(\"Модель успешно загружена.\")\n",
    "except Exception as e:\n",
    "    print(f\"Ошибка при загрузке модели: {e}\")\n",
    "\n",
    "# Define transformation for the input symbols\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.Resize((28, 28)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))  # Нормализация, при необходимости измените\n",
    "])\n",
    "\n",
    "# Main processing function\n",
    "def process_plate(plate):\n",
    "    \"\"\"\n",
    "    Processes a single license plate image to extract individual symbols based on masks.\n",
    "\n",
    "    :param plate: BGR image of the cropped license plate.\n",
    "    :return: Tuple containing the processed grayscale image and list of extracted symbols.\n",
    "    \"\"\"\n",
    "    # Convert to grayscale\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    # Resize for consistent processing (adjust fx and fy as needed)\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    img_h, img_w = resized_plate.shape\n",
    "\n",
    "    # Apply Gaussian blur to reduce noise\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "\n",
    "    # Extract symbols using both mask configurations\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\n",
    "\n",
    "    # Determine which mask to use based on the presence of the third region digit\n",
    "    # Check if the 9th symbol (Region Digit 3) is present\n",
    "    if len(symbols_9) == 9 and is_symbol_present(symbols_9[-1]):\n",
    "        characters = symbols_9\n",
    "    else:\n",
    "        characters = symbols_8\n",
    "\n",
    "    return blurred, characters\n",
    "\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "    # Ensure the plate image is valid\n",
    "    if plate is None:\n",
    "        print(f\"Plate {plate_index + 1} could not be loaded.\")\n",
    "        continue\n",
    "\n",
    "    # Process the plate to extract symbols\n",
    "    processed_plate, characters = process_plate(plate)\n",
    "\n",
    "    # Predict symbols using the model\n",
    "    predictions = []\n",
    "    for symbol_img in characters:\n",
    "        # Предобработка изображения\n",
    "        input_tensor = transform(symbol_img).unsqueeze(0).to(DEVICE)  # [1, 1, 28, 28]\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            outputs = model(input_tensor)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            predicted_label = CLASSES[predicted.item()]\n",
    "            predictions.append(predicted_label)\n",
    "    \n",
    "    # Display the results\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # Original Plate\n",
    "    plt.subplot(1, 4, 1)\n",
    "    plt.title(\"Original Plate\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Processed Plate (Grayscale & Blurred)\n",
    "    plt.subplot(1, 4, 2)\n",
    "    plt.title(\"Processed Plate\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Extracted Symbols\n",
    "    plt.subplot(1, 4, 3)\n",
    "    plt.title(\"Extracted Symbols\")\n",
    "    if characters:\n",
    "        # Arrange symbols in a horizontal grid\n",
    "        symbols_grid = np.hstack(characters)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Recognized Symbols with Predictions\n",
    "    plt.subplot(1, 4, 4)\n",
    "    plt.title(\"Recognized Symbols\")\n",
    "    if characters and predictions:\n",
    "        # Создадим холст для отображения символов с метками\n",
    "        symbol_images = [resize_with_padding(sym, size=28, padding=0, bg_color=255) for sym in characters]\n",
    "        num_symbols = len(symbol_images)\n",
    "        # Создадим пустой холст для размещения символов и меток\n",
    "        canvas = np.ones((28 + 10, 28 * num_symbols + 10), dtype=np.uint8) * 255  # 10 пикселей отступа\n",
    "        for i, (sym, pred) in enumerate(zip(symbol_images, predictions)):\n",
    "            canvas[0:28, i*28 + 5:(i+1)*28 + 5] = sym  # Расположим символы с отступом\n",
    "            # Добавим текст метки под символом\n",
    "            plt.text(i*28 + 14, 28 + 5, pred, ha='center', va='center', fontsize=8, color='blue')\n",
    "        plt.imshow(canvas, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.suptitle(f\"License Plate Processing and Recognition Result {plate_index + 1}\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Модель успешно загружена.\n",
      "Первая - T, Итоговая - T\n",
      "Первая - 8, Итоговая - 8\n",
      "Первая - 2, Итоговая - 2\n",
      "Первая - 9, Итоговая - 9\n",
      "Первая - M, Итоговая - M\n",
      "Первая - K, Итоговая - K\n",
      "Первая - 9, Итоговая - 9\n",
      "Первая - 7, Итоговая - 7\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 122724 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Первая - A, Итоговая - A\n",
      "Первая - 0, Итоговая - 0\n",
      "Первая - 2, Итоговая - 2\n",
      "Первая - 3, Итоговая - 3\n",
      "Первая - T, Итоговая - T\n",
      "Первая - Y, Итоговая - Y\n",
      "Первая - 9, Итоговая - 9\n",
      "Первая - 7, Итоговая - 7\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 131680 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Первая - K, Итоговая - K\n",
      "Первая - 2, Итоговая - 2\n",
      "Первая - 6, Итоговая - 6\n",
      "Первая - 3, Итоговая - 3\n",
      "Первая - C, Итоговая - C\n",
      "Первая - O, Итоговая - O\n",
      "Первая - 9, Итоговая - 9\n",
      "Первая - 7, Итоговая - 7\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 126644 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.Resize((28, 28)),  \n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.25,), (0.25,))  \n",
    "])\n",
    "\n",
    "def extract_symbols(image, masks):\n",
    "    \"\"\"\n",
    "    Извлекает символы из изображения на основе предоставленных масок.\n",
    "\n",
    "    :param image: Градация серого изображения номерного знака.\n",
    "    :param masks: Список кортежей с определениями масок как (x_frac, y_frac, w_frac, h_frac).\n",
    "    :return: Список извлечённых символов (numpy arrays).\n",
    "    \"\"\"\n",
    "    img_h, img_w = image.shape\n",
    "    symbols = []\n",
    "    for bbox in masks:\n",
    "        x_frac, y_frac, w_frac, h_frac = bbox\n",
    "        x = int(x_frac * img_w)\n",
    "        y = int(y_frac * img_h)\n",
    "        w = int(w_frac * img_w)\n",
    "        h = int(h_frac * img_h)\n",
    "        \n",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        \n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            continue\n",
    "\n",
    "        symbols.append(symbol)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    \"\"\"\n",
    "    Определяет, присутствует ли символ на основе средней интенсивности пикселей.\n",
    "\n",
    "    :param symbol: Градация серого изображения символа.\n",
    "    :param threshold: Порог интенсивности для определения наличия символа.\n",
    "    :return: Булево значение, указывающее на наличие символа.\n",
    "    \"\"\"\n",
    "    return np.mean(symbol) < threshold  \n",
    "\n",
    "class ImprovedCNN(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(ImprovedCNN, self).__init__()\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "            nn.SiLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(32, 64, kernel_size=3, padding=1),\n",
    "            nn.SiLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "\n",
    "            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n",
    "            nn.SiLU(),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.AdaptiveAvgPool2d((1, 1))  # Глобальный Average Pooling\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Linear(128, 256),\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(256, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv_layers(x)\n",
    "        x = x.view(x.size(0), -1)  # Разворачиваем перед линейным слоем\n",
    "        x = self.fc_layers(x)\n",
    "        return x\n",
    "model = ImprovedCNN(len(CLASSES)).to(DEVICE)\n",
    "SAVE_PATH = \"best_model_9.pth\"\n",
    "try:\n",
    "    model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "    model.eval()\n",
    "    print(\"Модель успешно загружена.\")\n",
    "except Exception as e:\n",
    "    print(f\"Ошибка при загрузке модели: {e}\")\n",
    "\n",
    "def process_plate(plate):\n",
    "    \"\"\"\n",
    "    Обрабатывает одно изображение номерного знака для извлечения отдельных символов на основе масок.\n",
    "\n",
    "    :param plate: BGR изображение обрезанного номерного знака.\n",
    "    :return: Кортеж, содержащий обработанное градационное изображение и список извлечённых символов.\n",
    "    \"\"\"\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    img_h, img_w = resized_plate.shape\n",
    "\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\n",
    "\n",
    "    if len(symbols_9) == 9 and is_symbol_present(symbols_9[-1]):\n",
    "        characters = symbols_9\n",
    "        expected_types = expected_types_9\n",
    "    else:\n",
    "        characters = symbols_8\n",
    "        expected_types = expected_types_8\n",
    "\n",
    "    return blurred, characters, expected_types\n",
    "\n",
    "    \n",
    "\n",
    "expected_types_8 = ['letter', 'digit', 'digit', 'digit', 'letter', 'letter', 'digit', 'digit']\n",
    "expected_types_9 = ['letter', 'digit', 'digit', 'digit', 'letter', 'letter', 'digit', 'digit', 'digit']\n",
    "\n",
    "digits = set(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'])\n",
    "letters = set([c for c in CLASSES if c.isalpha()])\n",
    "\n",
    "mask_8 = [\n",
    "    (0.03, 0.05, 0.12, 0.9),  # Символ 1: Буква\n",
    "    (0.16, 0.05, 0.12, 0.9),  # Символ 2: Цифра\n",
    "    (0.28, 0.05, 0.12, 0.9),  # Символ 3: Цифра\n",
    "    (0.40, 0.05, 0.12, 0.9),  # Символ 4: Цифра\n",
    "    (0.52, 0.05, 0.12, 0.9),  # Символ 5: Буква\n",
    "    (0.64, 0.05, 0.12, 0.9),  # Символ 6: Буква\n",
    "    (0.79, 0.07, 0.10, 0.60),  # Символ 7: Региональная цифра 1\n",
    "    (0.89, 0.07, 0.10, 0.60),  # Символ 8: Региональная цифра 2 \n",
    "]\n",
    "\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),  # Symbol 1: Letter\n",
    "    (0.17, 0.2, 0.1, 0.6),  # Symbol 2: Digit\n",
    "    (0.29, 0.2, 0.1, 0.6),  # Symbol 3: Digit\n",
    "    (0.41, 0.2, 0.1, 0.6),  # Symbol 4: Digit\n",
    "    (0.53, 0.2, 0.1, 0.6),  # Symbol 5: Letter\n",
    "    (0.65, 0.2, 0.1, 0.6),  # Symbol 6: Letter\n",
    "    (0.77, 0.2, 0.1, 0.6),  # Symbol 7: Region Digit 1\n",
    "    (0.89, 0.2, 0.1, 0.6),  # Symbol 8: Region Digit 2\n",
    "    (1.01, 0.2, 0.1, 0.6)   # Symbol 9: Region Digit 3\n",
    "]\n",
    "alls = []\n",
    "alls_t = []\n",
    "\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "\n",
    "    if plate is None:\n",
    "        print(f\"Plate {plate_index + 1} could not be loaded.\")\n",
    "        continue\n",
    "\n",
    "    processed_plate, characters, expected_types = process_plate(plate)\n",
    "\n",
    "    predictions = []\n",
    "    transformed_symbols = []\n",
    "    input_tensors = []\n",
    "    for idx, (symbol_img, expected_type) in enumerate(zip(characters, expected_types)):\n",
    "\n",
    "        input_tensor = transform(symbol_img).unsqueeze(0).to(DEVICE)\n",
    "        input_tensors.append(input_tensor)\n",
    "\n",
    "        transformed_image = transform(symbol_img).squeeze().cpu().numpy()\n",
    "        transformed_symbols.append(transformed_image)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(input_tensor)\n",
    "            probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
    "            probs, indices = torch.sort(probabilities, descending=True)\n",
    "\n",
    "            predicted_labels = [CLASSES[i] for i in indices[0]]\n",
    "            if expected_type == 'letter':\n",
    "\n",
    "                if predicted_labels[0] == '0':\n",
    "                    predicted_labels[0] = 'O'\n",
    "                elif predicted_labels[0] == '8':\n",
    "                    predicted_labels[0] = 'B'\n",
    "\n",
    "                for label in predicted_labels:\n",
    "                    if label in letters:\n",
    "                        predicted_label = label\n",
    "                        break\n",
    "                else:\n",
    "                    predicted_label = predicted_labels[0]\n",
    "\n",
    "            elif expected_type == 'digit':\n",
    "\n",
    "                if predicted_labels[0] == 'O':\n",
    "                    predicted_labels[0] = '0'\n",
    "                elif predicted_labels[0] == 'B':\n",
    "                    predicted_labels[0] = '8'\n",
    "\n",
    "                for label in predicted_labels:\n",
    "                    if label in digits:\n",
    "                        predicted_label = label\n",
    "                        break\n",
    "                else:\n",
    "                    predicted_label = predicted_labels[0]\n",
    "                \n",
    "\n",
    "            else:\n",
    "                predicted_label = predicted_labels[0]\n",
    "            print(f'Первая - {predicted_labels[0]}, Итоговая - {predicted_label}')\n",
    "            predictions.append(predicted_label)\n",
    "    alls.append(transformed_symbols)\n",
    "    alls_t.append(input_tensors)\n",
    "\n",
    "    plt.figure(figsize=(20, 5))\n",
    "\n",
    "\n",
    "    plt.subplot(1, 4, 1)\n",
    "    plt.title(\"Original Plate\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 2)\n",
    "    plt.title(\"Processed Plate\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 3)\n",
    "    plt.title(\"Extracted Symbols\")\n",
    "    if transformed_symbols:\n",
    "        symbols_grid = np.hstack(transformed_symbols)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 4)\n",
    "    plt.title(\"Recognized Symbols\")\n",
    "    if transformed_symbols and predictions:\n",
    "\n",
    "        num_symbols = len(transformed_symbols)\n",
    "        canvas = np.ones((28 + 10, 28 * num_symbols + 10), dtype=np.uint8) * 255 \n",
    "        for i, (sym, pred) in enumerate(zip(transformed_symbols, predictions)):\n",
    "            canvas[0:28, i*28 + 5:(i+1)*28 + 5] = sym \n",
    "            plt.text(i*28 + 14, 28 + 5, pred, ha='center', va='center', fontsize=8, color='blue')\n",
    "        plt.imshow(canvas, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "    \n",
    "\n",
    "    plt.suptitle(f\"License Plate Processing and Recognition Result {plate_index + 1}\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\leonk\\AppData\\Local\\Temp\\ipykernel_17656\\332507320.py:72: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Модель успешно загружена.\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'symbols_9' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[51], line 141\u001b[0m\n\u001b[0;32m    138\u001b[0m processed_plate, characters \u001b[38;5;241m=\u001b[39m process_plate(plate)\n\u001b[0;32m    140\u001b[0m \u001b[38;5;66;03m# Выбор соответствующего списка ожидаемых типов\u001b[39;00m\n\u001b[1;32m--> 141\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m characters \u001b[38;5;241m==\u001b[39m \u001b[43msymbols_9\u001b[49m:\n\u001b[0;32m    142\u001b[0m     expected_types \u001b[38;5;241m=\u001b[39m expected_types_9\n\u001b[0;32m    143\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "\u001b[1;31mNameError\u001b[0m: name 'symbols_9' is not defined"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "\n",
    "# Предполагается, что CLASSES и DEVICE уже определены\n",
    "# Например:\n",
    "# CLASSES = ['A', 'B', 'C', ..., '0', '1', ..., '9']\n",
    "# DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Разделение классов на буквы и цифры\n",
    "letters = [cls for cls in CLASSES if cls.isalpha()]\n",
    "digits = [cls for cls in CLASSES if cls.isdigit()]\n",
    "\n",
    "# Создание словарей для быстрого доступа к индексам\n",
    "letter_indices = {cls: idx for idx, cls in enumerate(CLASSES) if cls.isalpha()}\n",
    "digit_indices = {cls: idx for idx, cls in enumerate(CLASSES) if cls.isdigit()}\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(),\n",
    "    transforms.Resize((28, 28)),  \n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.25,), (0.25,))  \n",
    "])\n",
    "\n",
    "def extract_symbols(image, masks):\n",
    "    \"\"\"\n",
    "    Извлекает символы из изображения на основе предоставленных масок.\n",
    "\n",
    "    :param image: Градация серого изображения номерного знака.\n",
    "    :param masks: Список кортежей с определениями масок как (x_frac, y_frac, w_frac, h_frac).\n",
    "    :return: Список извлечённых символов (numpy arrays).\n",
    "    \"\"\"\n",
    "    img_h, img_w = image.shape\n",
    "    symbols = []\n",
    "    for bbox in masks:\n",
    "        x_frac, y_frac, w_frac, h_frac = bbox\n",
    "        x = int(x_frac * img_w)\n",
    "        y = int(y_frac * img_h)\n",
    "        w = int(w_frac * img_w)\n",
    "        h = int(h_frac * img_h)\n",
    "        \n",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        \n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            continue\n",
    "\n",
    "        symbols.append(symbol)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    \"\"\"\n",
    "    Определяет, присутствует ли символ на основе средней интенсивности пикселей.\n",
    "\n",
    "    :param symbol: Градация серого изображения символа.\n",
    "    :param threshold: Порог интенсивности для определения наличия символа.\n",
    "    :return: Булево значение, указывающее на наличие символа.\n",
    "    \"\"\"\n",
    "    return np.mean(symbol) < threshold  \n",
    "\n",
    "# Определение модели (предполагается, что SimpleCNN определен)\n",
    "model = SimpleCNN(len(CLASSES)).to(DEVICE)\n",
    "SAVE_PATH = \"best_model_8.pth\"\n",
    "try:\n",
    "    model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "    model.eval()\n",
    "    print(\"Модель успешно загружена.\")\n",
    "except Exception as e:\n",
    "    print(f\"Ошибка при загрузке модели: {e}\")\n",
    "\n",
    "def process_plate(plate):\n",
    "    \"\"\"\n",
    "    Обрабатывает одно изображение номерного знака для извлечения отдельных символов на основе масок.\n",
    "\n",
    "    :param plate: BGR изображение обрезанного номерного знака.\n",
    "    :return: Кортеж, содержащий обработанное градационное изображение и список извлечённых символов.\n",
    "    \"\"\"\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    img_h, img_w = resized_plate.shape\n",
    "\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\n",
    "\n",
    "    if len(symbols_9) == 9 and is_symbol_present(symbols_9[-1]):\n",
    "        characters = symbols_9\n",
    "    else:\n",
    "        characters = symbols_8\n",
    "\n",
    "    return blurred, characters\n",
    "\n",
    "mask_8 = [\n",
    "    (0.03, 0.05, 0.12, 0.9),  # Символ 1: Буква\n",
    "    (0.16, 0.05, 0.12, 0.9),  # Символ 2: Цифра\n",
    "    (0.28, 0.05, 0.12, 0.9),  # Символ 3: Цифра\n",
    "    (0.40, 0.05, 0.12, 0.9),  # Символ 4: Цифра\n",
    "    (0.52, 0.05, 0.12, 0.9),  # Символ 5: Буква\n",
    "    (0.64, 0.05, 0.12, 0.9),  # Символ 6: Буква\n",
    "    (0.80, 0.05, 0.11, 0.6),  # Символ 7: Региональная цифра 1\n",
    "    (0.90, 0.05, 0.11, 0.6),  # Символ 8: Региональная цифра 2 \n",
    "]\n",
    "\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),  # Symbol 1: Letter\n",
    "    (0.17, 0.2, 0.1, 0.6),  # Symbol 2: Digit\n",
    "    (0.29, 0.2, 0.1, 0.6),  # Symbol 3: Digit\n",
    "    (0.41, 0.2, 0.1, 0.6),  # Symbol 4: Digit\n",
    "    (0.53, 0.2, 0.1, 0.6),  # Symbol 5: Letter\n",
    "    (0.65, 0.2, 0.1, 0.6),  # Symbol 6: Letter\n",
    "    (0.77, 0.2, 0.1, 0.6),  # Symbol 7: Region Digit 1\n",
    "    (0.89, 0.2, 0.1, 0.6),  # Symbol 8: Region Digit 2\n",
    "    (1.01, 0.2, 0.1, 0.6)   # Symbol 9: Region Digit 3\n",
    "]\n",
    "\n",
    "# Определение ожидаемых типов символов для mask_8 и mask_9\n",
    "expected_types_8 = ['letter', 'digit', 'digit', 'digit', 'letter', 'letter', 'digit', 'digit']\n",
    "expected_types_9 = ['letter', 'digit', 'digit', 'digit', 'letter', 'letter', 'digit', 'digit', 'digit']\n",
    "\n",
    "alls = []\n",
    "alls_t = []\n",
    "\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "\n",
    "    if plate is None:\n",
    "        print(f\"Plate {plate_index + 1} could not be loaded.\")\n",
    "        continue\n",
    "\n",
    "    processed_plate, characters = process_plate(plate)\n",
    "\n",
    "    # Выбор соответствующего списка ожидаемых типов\n",
    "    if characters == symbols_9:\n",
    "        expected_types = expected_types_9\n",
    "    else:\n",
    "        expected_types = expected_types_8\n",
    "\n",
    "    predictions = []\n",
    "    transformed_symbols = [] \n",
    "    input_tensors = []\n",
    "    for idx, symbol_img in enumerate(characters):\n",
    "        input_tensor = transform(symbol_img).unsqueeze(0).to(DEVICE) \n",
    "        input_tensors.append(input_tensor)\n",
    "\n",
    "        transformed_image = transform(symbol_img).squeeze().cpu().numpy()  \n",
    "        transformed_symbols.append(transformed_image)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(input_tensor)  # Предполагается, что модель возвращает сырые логиты\n",
    "\n",
    "            # Применяем softmax для получения вероятностей\n",
    "            probabilities = F.softmax(outputs, dim=1).squeeze()\n",
    "\n",
    "            # Определение ожидаемого типа символа\n",
    "            expected_type = expected_types[idx]\n",
    "\n",
    "            if expected_type == 'letter':\n",
    "                # Получение индексов классов-букв\n",
    "                class_indices = [CLASSES.index(cls) for cls in letters]\n",
    "            elif expected_type == 'digit':\n",
    "                # Получение индексов классов-цифр\n",
    "                class_indices = [CLASSES.index(cls) for cls in digits]\n",
    "            else:\n",
    "                # Если тип не определен, использовать все классы\n",
    "                class_indices = list(range(len(CLASSES)))\n",
    "\n",
    "            # Извлечение вероятностей для соответствующих классов\n",
    "            if len(class_indices) == 0:\n",
    "                print(f\"Неизвестный тип символа: {expected_type}. Используются все классы.\")\n",
    "                class_indices = list(range(len(CLASSES)))\n",
    "\n",
    "            filtered_probs = probabilities[class_indices]\n",
    "            if filtered_probs.numel() == 0:\n",
    "                print(f\"Нет классов для типа: {expected_type}.\")\n",
    "                predicted_label = \"Unknown\"\n",
    "            else:\n",
    "                # Получение индекса максимальной вероятности среди отфильтрованных\n",
    "                max_prob, max_idx = torch.max(filtered_probs, dim=0)\n",
    "                predicted_label = CLASSES[class_indices[max_idx.item()]]\n",
    "\n",
    "            predictions.append(predicted_label)\n",
    "\n",
    "    alls.append(transformed_symbols)\n",
    "    alls_t.append(input_tensors)\n",
    "\n",
    "    plt.figure(figsize=(20, 5))\n",
    "\n",
    "    plt.subplot(1, 4, 1)\n",
    "    plt.title(\"Original Plate\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 2)\n",
    "    plt.title(\"Processed Plate\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 3)\n",
    "    plt.title(\"Extracted Symbols\")\n",
    "    if transformed_symbols:\n",
    "        symbols_grid = np.hstack(transformed_symbols)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(1, 4, 4)\n",
    "    plt.title(\"Recognized Symbols\")\n",
    "    if transformed_symbols and predictions:\n",
    "\n",
    "        num_symbols = len(transformed_symbols)\n",
    "        canvas = np.ones((28 + 20, 28 * num_symbols + 10), dtype=np.uint8) * 255 \n",
    "        for i, (sym, pred) in enumerate(zip(transformed_symbols, predictions)):\n",
    "            canvas[0:28, i*28 + 5:(i+1)*28 + 5] = sym \n",
    "            plt.text(i*28 + 14, 28 + 10, pred, ha='center', va='center', fontsize=12, color='blue')\n",
    "        plt.imshow(canvas, cmap='gray')\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'No Symbols Detected', ha='center', va='center', fontsize=12)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.suptitle(f\"License Plate Processing and Recognition Result {plate_index + 1}\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9600 characters]",
      "text/plain": [
       "<Figure size 200x200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def display_symbol(symbol_array, is_grayscale=True, title=\"Символ\"):\n",
    "    \"\"\"\n",
    "    Отображает символ из numpy массива.\n",
    "\n",
    "    :param symbol_array: Numpy массив, представляющий изображение символа.\n",
    "    :param is_grayscale: Булево значение, указывающее, является ли изображение чёрно-белым.\n",
    "    :param title: Заголовок для графика.\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(2, 2))  # Устанавливаем размер фигуры\n",
    "\n",
    "    if is_grayscale:\n",
    "        # Для чёрно-белых изображений используем cmap='gray'\n",
    "        plt.imshow(symbol_array, cmap='gray', interpolation='nearest')\n",
    "    else:\n",
    "        # Для цветных изображений предполагаем, что массив имеет форму (H, W, 3)\n",
    "        plt.imshow(symbol_array, interpolation='nearest')\n",
    "\n",
    "    plt.title(title)\n",
    "    plt.axis('off')  # Отключаем оси\n",
    "    plt.show()\n",
    "\n",
    "# Пример использования\n",
    "\n",
    "# Предполагается, что transformed_symbols уже определён и содержит хотя бы один символ\n",
    "# Например:\n",
    "# transformed_symbols = [numpy_array1, numpy_array2, ...]\n",
    "\n",
    "# Проверяем, что transformed_symbols не пуст\n",
    "try:\n",
    "    # Извлекаем первый символ\n",
    "    first_symbol = alls[0][1]\n",
    "    \n",
    "    # Определяем, является ли изображение чёрно-белым или цветным\n",
    "    if len(first_symbol.shape) == 2:\n",
    "        is_grayscale = True\n",
    "    elif len(first_symbol.shape) == 3 and first_symbol.shape[2] == 3:\n",
    "        is_grayscale = False\n",
    "    else:\n",
    "        raise ValueError(\"Неподдерживаемый формат изображения.\")\n",
    "    \n",
    "    # Отображаем первый символ\n",
    "    display_symbol(first_symbol, is_grayscale=is_grayscale, title=\"Первый извлечённый символ\")\n",
    "    \n",
    "except IndexError:\n",
    "    print(\"Список transformed_symbols пуст. Убедитесь, что символы были успешно извлечены.\")\n",
    "except Exception as e:\n",
    "    print(f\"Произошла ошибка при отображении символа: {e}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'LETTERS' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[32], line 7\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[0;32m      6\u001b[0m \u001b[38;5;66;03m# Инициализация модели\u001b[39;00m\n\u001b[1;32m----> 7\u001b[0m model_letters \u001b[38;5;241m=\u001b[39m SimpleCNN(\u001b[38;5;28mlen\u001b[39m(\u001b[43mLETTERS\u001b[49m))\u001b[38;5;241m.\u001b[39mto(DEVICE)  \u001b[38;5;66;03m# Инициализация модели для букв\u001b[39;00m\n\u001b[0;32m      8\u001b[0m model_letters\u001b[38;5;241m.\u001b[39mload_state_dict(torch\u001b[38;5;241m.\u001b[39mload(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbest_model_letters.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m, map_location\u001b[38;5;241m=\u001b[39mDEVICE))  \u001b[38;5;66;03m# Загрузка весов\u001b[39;00m\n\u001b[0;32m      9\u001b[0m model_letters\u001b[38;5;241m.\u001b[39meval()\n",
      "\u001b[1;31mNameError\u001b[0m: name 'LETTERS' is not defined"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import torch\n",
    "\n",
    "# Инициализация модели\n",
    "model_letters = SimpleCNN(len(LETTERS)).to(DEVICE)  # Инициализация модели для букв\n",
    "model_letters.load_state_dict(torch.load(\"best_model_letters.pth\", map_location=DEVICE))  # Загрузка весов\n",
    "model_letters.eval()\n",
    "\n",
    "model_digits = SimpleCNN(len(DIGITS)).to(DEVICE)  # Инициализация модели для цифр\n",
    "model_digits.load_state_dict(torch.load(\"best_model_digits.pth\", map_location=DEVICE))  # Загрузка весов\n",
    "model_digits.eval()\n",
    "\n",
    "def process_and_recognize_plate_with_processing(plate, model_digits, model_letters, transform, device=\"cpu\", padding=5):\n",
    "    \"\"\"\n",
    "    Обрабатывает пластину, распознаёт символы с дополнительной обработкой и возвращает текстовую строку.\n",
    "    \n",
    "    :param plate: Изображение пластины.\n",
    "    :param model_digits: Модель для распознавания цифр.\n",
    "    :param model_letters: Модель для распознавания букв.\n",
    "    :param transform: Преобразование изображения для модели.\n",
    "    :param device: Устройство для вычислений (\"cpu\" или \"cuda\").\n",
    "    :param padding: Отступ для символов.\n",
    "    :return: Обработанная пластина, список символов, распознанная строка.\n",
    "    \"\"\"\n",
    "    # Обработка пластины\n",
    "    processed_plate, characters = process_plate(plate)\n",
    "    \n",
    "    # Распознавание символов\n",
    "    recognized_text = \"\"\n",
    "    processed_characters = []\n",
    "    for i, char in enumerate(characters):\n",
    "        # Выбор модели для символа\n",
    "        if i in {0, 4, 5}:  # 1-й, 5-й и 6-й символы (нумерация с 0)\n",
    "            model = model_letters\n",
    "        else:\n",
    "            model = model_digits\n",
    "        \n",
    "        # Распознавание символа\n",
    "        recognized_char, processed_char = predict_character_with_processing(\n",
    "            model, char, transform, device, padding=padding\n",
    "        )\n",
    "        recognized_text += recognized_char\n",
    "        processed_characters.append(processed_char)\n",
    "    \n",
    "    return processed_plate, processed_characters, recognized_text\n",
    "\n",
    "# Пример использования с массивом processed_plates\n",
    "all_recognized_texts = []\n",
    "for plate_index, plate in enumerate(processed_plates):\n",
    "    # Обработка пластины и распознавание символов\n",
    "    processed_plate, processed_characters, recognized_text = process_and_recognize_plate_with_processing(\n",
    "        plate, model_digits, model_letters, transform, device=DEVICE, padding=5\n",
    "    )\n",
    "\n",
    "    # Сохранение распознанного текста\n",
    "    all_recognized_texts.append(recognized_text)\n",
    "\n",
    "    # Визуализация результатов\n",
    "    plt.figure(figsize=(15, 5))\n",
    "\n",
    "    # Исходная пластина\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.title(\"Исходная картинка\")\n",
    "    plt.imshow(cv2.cvtColor(plate, cv2.COLOR_BGR2RGB))  # Преобразуем для корректного отображения\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Обработанная пластина\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.title(\"Обработанная картинка\")\n",
    "    plt.imshow(processed_plate, cmap='gray')\n",
    "    plt.axis('off')\n",
    "\n",
    "    # Выделенные символы и итоговый текст\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.title(\"Выделенные символы\")\n",
    "    if processed_characters:\n",
    "        symbols_grid = np.hstack(processed_characters)\n",
    "        plt.imshow(symbols_grid, cmap='gray')\n",
    "        plt.xlabel(f\"Распознанный текст: {recognized_text}\")\n",
    "    else:\n",
    "        plt.text(0.5, 0.5, 'Нет символов', ha='center', va='center')\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.suptitle(f\"Результат обработки пластинки {plate_index + 1} - {recognized_text}\")\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20148 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21220 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20292 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20720 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19416 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20720 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20440 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20284 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Классы символов\n",
    "#CLASSES = \"DABEKMHOPCTYX0123456789\"\n",
    "#NUM_CLASSES = len(CLASSES)\n",
    "\n",
    "# Загрузка модели и весов\n",
    "#DEVICE = \"cpu\"  # Или \"cuda\", если доступен GPU\n",
    "#model = SimpleCNN(NUM_CLASSES).to(DEVICE)\n",
    "#model.load_state_dict(torch.load(\"best_model.pth\", map_location=DEVICE))  # Загрузка весов\n",
    "#model.eval()\n",
    "\n",
    "# Преобразование для входных данных\n",
    "transform = transforms.Compose([\n",
    "     # Изменение яркости (0.5 - уменьшение, >1 - увеличение)\n",
    "    transforms.ToTensor(),\n",
    "    #transforms.ColorJitter(brightness=0.5),\n",
    "    #transforms.Normalize((0.5,), (0.5,))  # Нормализация значений пикселей\n",
    "])\n",
    "\n",
    "# Функция для получения топ-3 предсказаний\n",
    "def get_top_predictions(output, top_k=3):\n",
    "    probabilities = torch.softmax(output, dim=1).squeeze()  # Преобразование в вероятности\n",
    "    top_probs, top_indices = torch.topk(probabilities, top_k)\n",
    "    return top_probs.tolist(), [CLASSES[idx] for idx in top_indices.tolist()]\n",
    "\n",
    "# Функция распознавания символов с выводом графиков\n",
    "def recognize_and_plot_characters(characters, model, device=\"cpu\"):\n",
    "    \"\"\"\n",
    "    Распознает символы из списка изображений characters и строит графики топ-3 предсказаний.\n",
    "\n",
    "    :param characters: Список символов (массивы numpy 28x28).\n",
    "    :param model: Загруженная модель для распознавания.\n",
    "    :param device: Устройство для вычислений (\"cpu\" или \"cuda\").\n",
    "    \"\"\"\n",
    "    for i, char in enumerate(characters):\n",
    "        # Преобразование символа в тензор\n",
    "        char_tensor = transform(char).unsqueeze(0).to(device)  # Добавляем batch размерности\n",
    "        \n",
    "        # Прогон через модель\n",
    "        with torch.no_grad():\n",
    "            output = model(char_tensor)\n",
    "            top_probs, top_classes = get_top_predictions(output)\n",
    "\n",
    "        # Вывод символа и графика топ-3 предсказаний\n",
    "        plt.figure(figsize=(8, 4))\n",
    "\n",
    "        # Символ\n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.imshow(char, cmap='gray')\n",
    "        plt.axis('off')\n",
    "        plt.title(f\"Символ {i + 1}\")\n",
    "\n",
    "        # Топ-3 предсказания\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.barh(top_classes, top_probs, color='blue')\n",
    "        plt.xlabel(\"Вероятность\")\n",
    "        plt.ylabel(\"Класс\")\n",
    "        plt.title(\"Топ-3 предсказания\")\n",
    "        plt.gca().invert_yaxis()  # Инвертируем ось, чтобы лучший класс был сверху\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "# Использование функции\n",
    "if 'characters' in locals() and len(transformed_symbols) > 0:\n",
    "    recognize_and_plot_characters(alls[1], model, device=DEVICE)\n",
    "else:\n",
    "    print(\"Список characters пуст.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Топ-10 предсказаний:\n",
      "1: B (вероятность: 0.3549)\n",
      "2: P (вероятность: 0.2256)\n",
      "3: 7 (вероятность: 0.0604)\n",
      "4: E (вероятность: 0.0508)\n",
      "5: 2 (вероятность: 0.0501)\n",
      "6: A (вероятность: 0.0425)\n",
      "7: M (вероятность: 0.0313)\n",
      "8: 8 (вероятность: 0.0289)\n",
      "9: 3 (вероятность: 0.0235)\n",
      "10: O (вероятность: 0.0192)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36360 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Предполагается, что CLASSES и DEVICE определены ранее\n",
    "# Например:\n",
    "# CLASSES = ['A', 'B', 'C', ..., 'Z', '0', '1', ..., '9']\n",
    "# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Путь к изображению\n",
    "image_path = \"img/7-viva4.jpg\"\n",
    "transform = transforms.Compose([\n",
    "        transforms.Resize((28, 28)),  # Изменяем размер до 28x28\n",
    "        transforms.ToTensor(),       # Преобразуем в тензор\n",
    "        transforms.Normalize(mean=[0.5], std=[0.5])  # Нормализация: среднее 0.5, стандартное отклонение 0.5\n",
    "    ])\n",
    "\n",
    "transform1 = transforms.Compose([\n",
    "        transforms.ToTensor(),       # Преобразуем в тензор\n",
    "        #transforms.Normalize(mean=[0.5], std=[0.5])  # Нормализация: среднее 0.5, стандартное отклонение 0.5\n",
    "    ])\n",
    "# Предобработка изображения для чёрно-белого изображения 28x28\n",
    "def preprocess_image(image_path):\n",
    "    # Загрузка изображения\n",
    "    image = Image.open(image_path).convert('L')  # Преобразование в чёрно-белое\n",
    "    # Применение трансформаций\n",
    "    image_tensor = transform(image)\n",
    "    \n",
    "    # Добавление batch dimension\n",
    "    image_tensor = image_tensor.unsqueeze(0)\n",
    "    return image_tensor\n",
    "\n",
    "# Загрузка и обработка изображения\n",
    "image_tensor = preprocess_image(image_path).to(DEVICE)\n",
    "char_tensor = transform1(transformed_symbols[3]).unsqueeze(0).to(DEVICE)\n",
    "\n",
    "# Загрузка модели (предполагается, что модель уже определена)\n",
    "# model = SimpleCNN(len(CLASSES)).to(DEVICE)\n",
    "# Загрузка весов модели\n",
    "# model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "# model.eval()\n",
    "\n",
    "# Убедитесь, что модель находится в режиме оценки\n",
    "model.eval()\n",
    "\n",
    "# Добавление вывода изображения и графика топ-10 предсказаний\n",
    "def plot_image_and_predictions(image_tensor, classes):\n",
    "    \"\"\"\n",
    "    Отображает изображение и график топ-10 предсказаний модели.\n",
    "    \n",
    "    :param image_tensor: Тензор изображения после предобработки.\n",
    "    :param top_probs: Топ-10 вероятностей предсказаний.\n",
    "    :param top_indices: Индексы топ-10 классов.\n",
    "    :param classes: Список названий классов.\n",
    "    \"\"\"\n",
    "    # Обратная нормализация для отображения изображения\n",
    "    # Если нормализация была (mean=0.5, std=0.5), то обратная:\n",
    "    # x = x * std + mean\n",
    "\n",
    "    # Предсказание\n",
    "    with torch.no_grad():\n",
    "        output = model(image_tensor)  # Получение выхода модели\n",
    "        probabilities = torch.softmax(output, dim=1)  # Преобразование выходов в вероятности\n",
    "        top10_probabilities, top10_indices = torch.topk(probabilities, 10, dim=1)  # Получение топ-10\n",
    "        \n",
    "        top_probs = top10_probabilities.squeeze(0)  # Удаление batch-дименшена\n",
    "        top_indices = top10_indices.squeeze(0)  # Удаление batch-дименшена\n",
    "\n",
    "    print(\"Топ-10 предсказаний:\")\n",
    "    for i in range(10):\n",
    "        predicted_class = top_indices[i].item()\n",
    "        probability = top_probs[i].item()\n",
    "        print(f\"{i + 1}: {CLASSES[predicted_class]} (вероятность: {probability:.4f})\")\n",
    "        \n",
    "    unnormalize = transforms.Normalize(\n",
    "        mean=[-0.5 / 0.5],\n",
    "        std=[1 / 0.5]\n",
    "    )\n",
    "    image_tensor_unnorm = unnormalize(image_tensor.squeeze(0)).cpu().numpy()\n",
    "    \n",
    "    # Переводим тензор в формат [H, W]\n",
    "    image_np = image_tensor_unnorm.squeeze(0)\n",
    "    image_np = np.clip(image_np, 0, 1)  # Ограничиваем значения для корректного отображения\n",
    "    \n",
    "    # Получение топ-10 классов и их вероятностей\n",
    "    top_classes = [classes[idx] for idx in top_indices]\n",
    "    top_probs = top_probs.cpu().numpy()\n",
    "    \n",
    "    # Создание фигуры с двумя подграфиками\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "    \n",
    "    # Отображение изображения\n",
    "    ax1.imshow(image_np, cmap='gray')\n",
    "    ax1.set_title(\"Предсказанное изображение\")\n",
    "    ax1.axis('off')\n",
    "    \n",
    "    # Отображение графика топ-10 предсказаний\n",
    "    ax2.barh(range(10), top_probs, color='skyblue')  # Инвертируем для отображения сверху вниз\n",
    "    ax2.set_yticks(range(10))\n",
    "    ax2.set_yticklabels(top_classes)\n",
    "    ax2.invert_yaxis()  # Самый высокий показатель сверху\n",
    "    ax2.set_xlabel('Вероятность')\n",
    "    ax2.set_title('Топ-10 предсказаний')\n",
    "    \n",
    "    # Добавление вероятностей на график\n",
    "    for i, v in enumerate(top_probs):\n",
    "        ax2.text(v + 0.01, i, f\"{v:.2f}\", color='blue', va='center')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Вызов функции для отображения\n",
    "plot_image_and_predictions(image_tensor, CLASSES)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37864 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Топ-10 предсказаний:\n",
      "1: Y (вероятность: 0.2801)\n",
      "2: 2 (вероятность: 0.1730)\n",
      "3: 9 (вероятность: 0.0611)\n",
      "4: M (вероятность: 0.0609)\n",
      "5: P (вероятность: 0.0556)\n",
      "6: X (вероятность: 0.0436)\n",
      "7: 8 (вероятность: 0.0414)\n",
      "8: H (вероятность: 0.0380)\n",
      "9: 0 (вероятность: 0.0329)\n",
      "10: 1 (вероятность: 0.0279)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34376 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import cv2 as cv\n",
    "\n",
    "# Предполагается, что CLASSES и DEVICE определены ранее\n",
    "# Например:\n",
    "# CLASSES = ['A', 'B', 'C', ..., 'Z', '0', '1', ..., '9']\n",
    "# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Путь к изображению\n",
    "image_path = \"img/7-bad1.png\"\n",
    "\n",
    "# Определение трансформаций\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((28, 28)),  # Изменяем размер до 28x28\n",
    "    transforms.ToTensor(),        # Преобразуем в тензор\n",
    "    transforms.Normalize(mean=[0.5], std=[0.5])  # Нормализация: среднее 0.5, стандартное отклонение 0.5\n",
    "])\n",
    "\n",
    "transform1 = transforms.Compose([\n",
    "    transforms.ToTensor(),        # Преобразуем в тензор\n",
    "    #transforms.Normalize(mean=[0.5], std=[0.5])  # Нормализация: среднее 0.5, стандартное отклонение 0.5\n",
    "])\n",
    "\n",
    "def preprocess_image(image_path):\n",
    "    # Шаг 1: Загрузка изображения с помощью OpenCV\n",
    "    src = cv.imread(image_path, cv.IMREAD_COLOR)\n",
    "    assert src is not None, f\"Ошибка при открытии изображения по пути: {image_path}\"\n",
    "\n",
    "    # Шаг 2: Преобразование в оттенки серого\n",
    "    if len(src.shape) != 2:\n",
    "        gray = cv.cvtColor(src, cv.COLOR_BGR2GRAY)\n",
    "    else:\n",
    "        gray = src\n",
    "\n",
    "    # Инверсия цветов\n",
    "    gray = cv.bitwise_not(gray)\n",
    "\n",
    "    # Бинаризация изображения\n",
    "    bw = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 15, -2)\n",
    "\n",
    "    # Извлечение горизонтальных и вертикальных линий\n",
    "    horizontal = np.copy(bw)\n",
    "    vertical = np.copy(bw)\n",
    "\n",
    "    # Обработка горизонтальных линий\n",
    "    horizontal_size = horizontal.shape[1] // 30\n",
    "    horizontalStructure = cv.getStructuringElement(cv.MORPH_RECT, (horizontal_size, 1))\n",
    "    horizontal = cv.erode(horizontal, horizontalStructure)\n",
    "    horizontal = cv.dilate(horizontal, horizontalStructure)\n",
    "\n",
    "    # Обработка вертикальных линий\n",
    "    vertical_size = vertical.shape[0] // 30\n",
    "    verticalStructure = cv.getStructuringElement(cv.MORPH_RECT, (1, vertical_size))\n",
    "    vertical = cv.erode(vertical, verticalStructure)\n",
    "    vertical = cv.dilate(vertical, verticalStructure)\n",
    "\n",
    "    # Инверсия вертикальных линий\n",
    "    vertical = cv.bitwise_not(vertical)\n",
    "\n",
    "    # Выделение краев\n",
    "    edges = cv.adaptiveThreshold(vertical, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 3, -2)\n",
    "    kernel = np.ones((2, 2), np.uint8)\n",
    "    edges = cv.dilate(edges, kernel)\n",
    "\n",
    "    # Сглаживание вертикальных линий\n",
    "    smooth = np.copy(vertical)\n",
    "    smooth = cv.blur(smooth, (2, 2))\n",
    "    rows, cols = np.where(edges != 0)\n",
    "    vertical[rows, cols] = smooth[rows, cols]\n",
    "\n",
    "    # Финальная обработка: использование сглаженных вертикальных линий\n",
    "    final_processed = vertical\n",
    "\n",
    "    # (Опционально) Визуализация этапов обработки\n",
    "    \n",
    "    titles = ['Original', 'Grayscale', 'Binary', 'Horizontal Lines', 'Vertical Lines', 'Smoothed Vertical']\n",
    "    images = [src, gray, bw, horizontal, vertical, smooth]\n",
    "\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    for i in range(len(images)):\n",
    "        plt.subplot(2, 3, i + 1)\n",
    "        plt.imshow(images[i], cmap='gray' if len(images[i].shape) == 2 else None)\n",
    "        plt.title(titles[i])\n",
    "        plt.axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "\n",
    "    # Преобразование обработанного изображения в формат PIL для применения трансформаций\n",
    "    final_image = Image.fromarray(final_processed)\n",
    "\n",
    "    # Применение трансформаций\n",
    "    image_tensor = transform(final_image)\n",
    "\n",
    "    # Добавление batch dimension\n",
    "    image_tensor = image_tensor.unsqueeze(0)\n",
    "\n",
    "    return image_tensor\n",
    "\n",
    "# Загрузка и обработка изображения\n",
    "image_tensor = preprocess_image(image_path).to(DEVICE)\n",
    "\n",
    "# Обработка символов (предполагается, что transformed_symbols определен)\n",
    "# char_tensor = transform1(transformed_symbols[3]).unsqueeze(0).to(DEVICE)\n",
    "\n",
    "# Загрузка модели (предполагается, что модель уже определена)\n",
    "# model = SimpleCNN(len(CLASSES)).to(DEVICE)\n",
    "# Загрузка весов модели\n",
    "# model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "# model.eval()\n",
    "\n",
    "# Убедитесь, что модель находится в режиме оценки\n",
    "model.eval()\n",
    "\n",
    "# Добавление вывода изображения и графика топ-10 предсказаний\n",
    "def plot_image_and_predictions(image_tensor, classes):\n",
    "    \"\"\"\n",
    "    Отображает изображение и график топ-10 предсказаний модели.\n",
    "    \n",
    "    :param image_tensor: Тензор изображения после предобработки.\n",
    "    :param classes: Список названий классов.\n",
    "    \"\"\"\n",
    "    # Предсказание\n",
    "    with torch.no_grad():\n",
    "        output = model(image_tensor)  # Получение выхода модели\n",
    "        probabilities = torch.softmax(output, dim=1)  # Преобразование выходов в вероятности\n",
    "        top10_probabilities, top10_indices = torch.topk(probabilities, 10, dim=1)  # Получение топ-10\n",
    "        \n",
    "        top_probs = top10_probabilities.squeeze(0)  # Удаление batch-дименшена\n",
    "        top_indices = top10_indices.squeeze(0)      # Удаление batch-дименшена\n",
    "\n",
    "    print(\"Топ-10 предсказаний:\")\n",
    "    for i in range(10):\n",
    "        predicted_class = top_indices[i].item()\n",
    "        probability = top_probs[i].item()\n",
    "        print(f\"{i + 1}: {CLASSES[predicted_class]} (вероятность: {probability:.4f})\")\n",
    "        \n",
    "    # Обратная нормализация для отображения изображения\n",
    "    unnormalize = transforms.Normalize(\n",
    "        mean=[-0.5 / 0.5],\n",
    "        std=[1 / 0.5]\n",
    "    )\n",
    "    image_tensor_unnorm = unnormalize(image_tensor.squeeze(0)).cpu().numpy()\n",
    "    \n",
    "    # Переводим тензор в формат [H, W]\n",
    "    image_np = image_tensor_unnorm.squeeze(0)\n",
    "    image_np = np.clip(image_np, 0, 1)  # Ограничиваем значения для корректного отображения\n",
    "    \n",
    "    # Получение топ-10 классов и их вероятностей\n",
    "    top_classes = [classes[idx] for idx in top_indices]\n",
    "    top_probs = top_probs.cpu().numpy()\n",
    "    \n",
    "    # Создание фигуры с двумя подграфиками\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "    \n",
    "    # Отображение изображения\n",
    "    ax1.imshow(image_np, cmap='gray')\n",
    "    ax1.set_title(\"Предсказанное изображение\")\n",
    "    ax1.axis('off')\n",
    "    \n",
    "    # Отображение графика топ-10 предсказаний\n",
    "    ax2.barh(range(10), top_probs, color='skyblue')  # Инвертируем для отображения сверху вниз\n",
    "    ax2.set_yticks(range(10))\n",
    "    ax2.set_yticklabels(top_classes)\n",
    "    ax2.invert_yaxis()  # Самый высокий показатель сверху\n",
    "    ax2.set_xlabel('Вероятность')\n",
    "    ax2.set_title('Топ-10 предсказаний')\n",
    "    \n",
    "    # Добавление вероятностей на график\n",
    "    for i, v in enumerate(top_probs):\n",
    "        ax2.text(v + 0.01, i, f\"{v:.2f}\", color='blue', va='center')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Вызов функции для отображения\n",
    "plot_image_and_predictions(image_tensor, CLASSES)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18512 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19148 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19320 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19076 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19864 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19844 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20464 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18040 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image, ImageOps\n",
    "import cv2\n",
    "\n",
    "# Классы символов\n",
    "CLASSES = \"ABEKMHOPCTYX0123456789\"\n",
    "NUM_CLASSES = len(CLASSES)\n",
    "\n",
    "# Загрузка модели и весов\n",
    "DEVICE = \"cpu\"  # Или \"cuda\", если доступен GPU\n",
    "model = SimpleCNN(NUM_CLASSES).to(DEVICE)\n",
    "model.load_state_dict(torch.load(\"best_model.pth\", map_location=DEVICE))  # Загрузка весов\n",
    "model.eval()\n",
    "\n",
    "# Преобразование для входных данных\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5,), (0.5,))  # Нормализация значений пикселей\n",
    "])\n",
    "\n",
    "# Функция для получения топ-3 предсказаний\n",
    "def get_top_predictions(output, top_k=3):\n",
    "    probabilities = torch.softmax(output, dim=1).squeeze()  # Преобразование в вероятности\n",
    "    top_probs, top_indices = torch.topk(probabilities, top_k)\n",
    "    return top_probs.tolist(), [CLASSES[idx] for idx in top_indices.tolist()]\n",
    "\n",
    "# Функция для добавления отступов\n",
    "def add_padding(img, padding=10):\n",
    "    \"\"\"\n",
    "    Добавляет отступы вокруг изображения.\n",
    "\n",
    "    :param img: Массив numpy с изображением символа.\n",
    "    :param padding: Количество пикселей для отступа.\n",
    "    :return: Изображение с отступами.\n",
    "    \"\"\"\n",
    "    pil_img = Image.fromarray(img)\n",
    "    padded_img = ImageOps.expand(pil_img, border=padding, fill=0)\n",
    "    return np.array(padded_img)\n",
    "\n",
    "# Функция для утончения линий\n",
    "def thin_lines(img, kernel_size=(2, 2), iterations=1):\n",
    "    \"\"\"\n",
    "    Утончает линии на изображении с помощью морфологической операции эрозии.\n",
    "\n",
    "    :param img: Массив numpy с изображением символа.\n",
    "    :param kernel_size: Размер ядра эрозии.\n",
    "    :param iterations: Количество итераций эрозии.\n",
    "    :return: Изображение с утонченными линиями.\n",
    "    \"\"\"\n",
    "    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size)\n",
    "    eroded_img = cv2.erode(img, kernel, iterations=iterations)\n",
    "    return eroded_img\n",
    "\n",
    "# Функция распознавания символов с выводом графиков\n",
    "# Функция для изменения размера изображения\n",
    "# Функция для изменения размера изображения\n",
    "def resize_to_model_input(img, target_size=(28, 28)):\n",
    "    \"\"\"\n",
    "    Изменяет размер изображения до требуемого для модели.\n",
    "\n",
    "    :param img: Массив numpy с изображением символа.\n",
    "    :param target_size: Размер, до которого нужно привести изображение.\n",
    "    :return: Изображение с измененным размером.\n",
    "    \"\"\"\n",
    "    pil_img = Image.fromarray(img)\n",
    "    resized_img = pil_img.resize(target_size, Image.Resampling.LANCZOS)  # Используем LANCZOS вместо ANTIALIAS\n",
    "    return np.array(resized_img)\n",
    "\n",
    "# Остальной код остается без изменений\n",
    "\n",
    "\n",
    "# Функция распознавания символов с исправлением размера\n",
    "def recognize_and_plot_characters(characters, model, device=\"cpu\", padding=5):\n",
    "    \"\"\"\n",
    "    Распознает символы из списка изображений characters и строит графики топ-3 предсказаний.\n",
    "\n",
    "    :param characters: Список символов (массивы numpy 28x28).\n",
    "    :param model: Загруженная модель для распознавания.\n",
    "    :param device: Устройство для вычислений (\"cpu\" или \"cuda\").\n",
    "    :param padding: Отступы вокруг символа.\n",
    "    \"\"\"\n",
    "    for i, char in enumerate(characters):\n",
    "        # Добавление отступов\n",
    "        char_padded = add_padding(char, padding=padding)\n",
    "\n",
    "        \n",
    "        char_selective = selective_dilation(char_padded, distance_threshold=4, iterations=1)\n",
    "        char_selective = thin_lines(char_selective, kernel_size=(2,2), iterations=1)\n",
    "        # Изменение размера для модели\n",
    "        char_resized = resize_to_model_input(char_selective, target_size=(28, 28))\n",
    "\n",
    "        # Преобразование символа в тензор\n",
    "        char_tensor = transform(char_resized).unsqueeze(0).to(device)  # Добавляем batch размерности\n",
    "        \n",
    "        # Прогон через модель\n",
    "        with torch.no_grad():\n",
    "            output = model(char_tensor)\n",
    "            top_probs, top_classes = get_top_predictions(output)\n",
    "\n",
    "        # Вывод символа и графика топ-3 предсказаний\n",
    "        plt.figure(figsize=(8, 4))\n",
    "\n",
    "        # Символ\n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.imshow(char_resized, cmap='gray')\n",
    "        plt.axis('off')\n",
    "        plt.title(f\"Символ {i + 1}\")\n",
    "\n",
    "        # Топ-3 предсказания\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.barh(top_classes, top_probs, color='blue')\n",
    "        plt.xlabel(\"Вероятность\")\n",
    "        plt.ylabel(\"Класс\")\n",
    "        plt.title(\"Топ-3 предсказания\")\n",
    "        plt.gca().invert_yaxis()  # Инвертируем ось, чтобы лучший класс был сверху\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "# Использование функции\n",
    "if 'characters' in locals() and len(characters) > 0:\n",
    "    recognize_and_plot_characters(characters, model, device=DEVICE, padding=5)\n",
    "else:\n",
    "    print(\"Список characters пуст.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
