{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "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": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1088772 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 986872 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 943732 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",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms\n",
    "\n",
    "# Конфигурация устройства\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Классы символов для модели OCR\n",
    "CLASSES = \"ABEKMHOPCTYX0123456789\"\n",
    "\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",
    "# Определение модели CNN\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",
    "# Загрузка предобученной модели\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",
    "except Exception as e:\n",
    "    print(f\"Ошибка при загрузке модели: {e}\")\n",
    "\n",
    "# Функции для распознавания номера\n",
    "def extract_symbols(image, masks):\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",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            continue\n",
    "        symbols.append(symbol)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    return np.mean(symbol) < threshold\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",
    "mask_8 = [\n",
    "    (0.03, 0.05, 0.12, 0.9),\n",
    "    (0.16, 0.05, 0.12, 0.9),\n",
    "    (0.28, 0.05, 0.12, 0.9),\n",
    "    (0.40, 0.05, 0.12, 0.9),\n",
    "    (0.52, 0.05, 0.12, 0.9),\n",
    "    (0.64, 0.05, 0.12, 0.9),\n",
    "    (0.79, 0.07, 0.10, 0.60),\n",
    "    (0.89, 0.07, 0.10, 0.60),\n",
    "]\n",
    "\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),\n",
    "    (0.17, 0.2, 0.1, 0.6),\n",
    "    (0.29, 0.2, 0.1, 0.6),\n",
    "    (0.41, 0.2, 0.1, 0.6),\n",
    "    (0.53, 0.2, 0.1, 0.6),\n",
    "    (0.65, 0.2, 0.1, 0.6),\n",
    "    (0.77, 0.2, 0.1, 0.6),\n",
    "    (0.89, 0.2, 0.1, 0.6),\n",
    "    (1.01, 0.2, 0.1, 0.6)\n",
    "]\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",
    "def process_plate(plate):\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\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",
    "    return blurred, characters, expected_types\n",
    "\n",
    "def recognize_plate(plate_image):\n",
    "    if plate_image is None:\n",
    "        return None\n",
    "    processed_plate, characters, expected_types = process_plate(plate_image)\n",
    "    predictions = []\n",
    "    for idx, (symbol_img, expected_type) in enumerate(zip(characters, expected_types)):\n",
    "        input_tensor = transform(symbol_img).unsqueeze(0).to(DEVICE)\n",
    "        with torch.no_grad():\n",
    "            outputs = model(input_tensor)\n",
    "            probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
    "            indices = torch.argmax(probabilities, dim=1)\n",
    "            predicted_label = CLASSES[indices[0]]\n",
    "            if expected_type == 'letter':\n",
    "                if predicted_label == '0':\n",
    "                    predicted_label = 'O'\n",
    "                elif predicted_label == '8':\n",
    "                    predicted_label = 'B'\n",
    "                if predicted_label not in letters:\n",
    "                    continue\n",
    "            elif expected_type == 'digit':\n",
    "                if predicted_label == 'O':\n",
    "                    predicted_label = '0'\n",
    "                elif predicted_label == 'B':\n",
    "                    predicted_label = '8'\n",
    "                if predicted_label not in digits:\n",
    "                    continue\n",
    "            predictions.append(predicted_label)\n",
    "    plate_number = ''.join(predictions)\n",
    "    return plate_number\n",
    "\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",
    "        # Распознавание номера\n",
    "        plate_number = recognize_plate(plate_image)\n",
    "    else:\n",
    "        plate_number = None\n",
    "\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",
    "        # Вывод распознанного номера\n",
    "        if plate_number:\n",
    "            placeholder_text = \"Number plate:\"\n",
    "            decoded_text = plate_number\n",
    "        else:\n",
    "            placeholder_text = \"Not found\"\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'Результат обработки: {image_path}')\n",
    "    plt.show()\n",
    "\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": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "0: 480x640 1 , 7.0ms\n",
      "Speed: 2.0ms preprocess, 7.0ms inference, 1.0ms postprocess per image at shape (1, 3, 480, 640)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\leonk\\AppData\\Local\\Temp\\ipykernel_6660\\3663620689.py:243: DeprecationWarning: `np.int0` is a deprecated alias for `np.intp`.  (Deprecated NumPy 1.24)\n",
      "  cv2.polylines(output_image, [np.int0(best_plate_box)], isClosed=True, color=(0, 255, 0), thickness=2)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 60256 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\leonk\\AppData\\Local\\Temp\\ipykernel_6660\\3663620689.py:255: DeprecationWarning: `np.int0` is a deprecated alias for `np.intp`.  (Deprecated NumPy 1.24)\n",
      "  box = np.int0(box)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 20400 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 1087372 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "0: 480x640 1 , 8.0ms\n",
      "Speed: 2.0ms preprocess, 8.0ms inference, 1.0ms postprocess per image at shape (1, 3, 480, 640)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57984 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22660 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 982092 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "0: 480x640 1 , 7.8ms\n",
      "Speed: 1.0ms preprocess, 7.8ms inference, 1.0ms postprocess per image at shape (1, 3, 480, 640)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 47844 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22152 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 938044 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "0: 448x640 1 , 10.6ms\n",
      "Speed: 2.0ms preprocess, 10.6ms inference, 1.0ms postprocess per image at shape (1, 3, 448, 640)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 286684 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 198004 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 893736 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",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms\n",
    "from ultralytics import YOLO\n",
    "\n",
    "# Configuration\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# OCR Classes\n",
    "CLASSES = \"ABEKMHOPCTYX0123456789\"\n",
    "\n",
    "# Image Transformations\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",
    "# Define the OCR Model\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))\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",
    "\n",
    "# Load the OCR Model\n",
    "ocr_model = ImprovedCNN(len(CLASSES)).to(DEVICE)\n",
    "SAVE_PATH = \"best_model_9.pth\"\n",
    "try:\n",
    "    ocr_model.load_state_dict(torch.load(SAVE_PATH, map_location=DEVICE))\n",
    "    ocr_model.eval()\n",
    "except Exception as e:\n",
    "    print(f\"Error loading OCR model: {e}\")\n",
    "\n",
    "# Load YOLO Model\n",
    "yolo_model = YOLO(\"yolo_plate.pt\")\n",
    "\n",
    "# OCR Functions (Same as before)\n",
    "def extract_symbols(image, masks):\n",
    "    # Function code remains the same\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",
    "        x_end = min(x + w, img_w)\n",
    "        y_end = min(y + h, img_h)\n",
    "        symbol = image[y:y_end, x:x_end]\n",
    "        if symbol.size == 0:\n",
    "            continue\n",
    "        symbols.append(symbol)\n",
    "    return symbols\n",
    "\n",
    "def is_symbol_present(symbol, threshold=250):\n",
    "    return np.mean(symbol) < threshold\n",
    "\n",
    "# Masks and Expected Types (Same as before)\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",
    "mask_8 = [\n",
    "    (0.03, 0.05, 0.12, 0.9),\n",
    "    (0.16, 0.05, 0.12, 0.9),\n",
    "    (0.28, 0.05, 0.12, 0.9),\n",
    "    (0.40, 0.05, 0.12, 0.9),\n",
    "    (0.52, 0.05, 0.12, 0.9),\n",
    "    (0.64, 0.05, 0.12, 0.9),\n",
    "    (0.79, 0.07, 0.10, 0.60),\n",
    "    (0.89, 0.07, 0.10, 0.60),\n",
    "]\n",
    "\n",
    "mask_9 = [\n",
    "    (0.05, 0.2, 0.1, 0.6),\n",
    "    (0.17, 0.2, 0.1, 0.6),\n",
    "    (0.29, 0.2, 0.1, 0.6),\n",
    "    (0.41, 0.2, 0.1, 0.6),\n",
    "    (0.53, 0.2, 0.1, 0.6),\n",
    "    (0.65, 0.2, 0.1, 0.6),\n",
    "    (0.77, 0.2, 0.1, 0.6),\n",
    "    (0.89, 0.2, 0.1, 0.6),\n",
    "    (1.01, 0.2, 0.1, 0.6)\n",
    "]\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",
    "def process_plate(plate):\n",
    "    # Function code remains the same\n",
    "    gray_plate = cv2.cvtColor(plate, cv2.COLOR_BGR2GRAY)\n",
    "    resized_plate = cv2.resize(gray_plate, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)\n",
    "    blurred = cv2.GaussianBlur(resized_plate, (3, 3), 0)\n",
    "    symbols_8 = extract_symbols(blurred, mask_8)\n",
    "    symbols_9 = extract_symbols(blurred, mask_9)\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",
    "    return blurred, characters, expected_types\n",
    "\n",
    "def recognize_plate(plate_image):\n",
    "    # Function code remains the same\n",
    "    if plate_image is None:\n",
    "        return None\n",
    "    processed_plate, characters, expected_types = process_plate(plate_image)\n",
    "    predictions = []\n",
    "    for idx, (symbol_img, expected_type) in enumerate(zip(characters, expected_types)):\n",
    "        input_tensor = transform(symbol_img).unsqueeze(0).to(DEVICE)\n",
    "        with torch.no_grad():\n",
    "            outputs = ocr_model(input_tensor)\n",
    "            probabilities = torch.nn.functional.softmax(outputs, dim=1)\n",
    "            indices = torch.argmax(probabilities, dim=1)\n",
    "            predicted_label = CLASSES[indices[0]]\n",
    "            if expected_type == 'letter':\n",
    "                if predicted_label == '0':\n",
    "                    predicted_label = 'O'\n",
    "                elif predicted_label == '8':\n",
    "                    predicted_label = 'B'\n",
    "                if predicted_label not in letters:\n",
    "                    continue\n",
    "            elif expected_type == 'digit':\n",
    "                if predicted_label == 'O':\n",
    "                    predicted_label = '0'\n",
    "                elif predicted_label == 'B':\n",
    "                    predicted_label = '8'\n",
    "                if predicted_label not in digits:\n",
    "                    continue\n",
    "            predictions.append(predicted_label)\n",
    "    plate_number = ''.join(predictions)\n",
    "    return plate_number\n",
    "\n",
    "# Perspective Transformation Functions\n",
    "def order_points(pts):\n",
    "    rect = np.zeros((4, 2), dtype=\"float32\")\n",
    "    pts = np.array(pts)\n",
    "    s = pts.sum(axis=1)\n",
    "    rect[0] = pts[np.argmin(s)]  # Top-left\n",
    "    rect[2] = pts[np.argmax(s)]  # Bottom-right\n",
    "    diff = np.diff(pts, axis=1)\n",
    "    rect[1] = pts[np.argmin(diff)]  # Top-right\n",
    "    rect[3] = pts[np.argmax(diff)]  # Bottom-left\n",
    "    return rect\n",
    "\n",
    "def four_point_transform(image, pts):\n",
    "    rect = order_points(pts)\n",
    "    (tl, tr, br, bl) = rect\n",
    "    widthA = np.hypot(br[0] - bl[0], br[1] - bl[1])\n",
    "    widthB = np.hypot(tr[0] - tl[0], tr[1] - tl[1])\n",
    "    maxWidth = max(int(widthA), int(widthB))\n",
    "    heightA = np.hypot(tr[0] - br[0], tr[1] - br[1])\n",
    "    heightB = np.hypot(tl[0] - bl[0], tl[1] - bl[1])\n",
    "    maxHeight = max(int(heightA), int(heightB))\n",
    "    dst = np.array([\n",
    "        [0, 0],\n",
    "        [maxWidth - 1, 0],\n",
    "        [maxWidth - 1, maxHeight - 1],\n",
    "        [0, maxHeight - 1]], dtype=\"float32\")\n",
    "    M = cv2.getPerspectiveTransform(rect, dst)\n",
    "    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))\n",
    "    return warped\n",
    "\n",
    "# Debug Function\n",
    "def show_debug_image(title, image):\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    if len(image.shape) == 2:  # grayscale\n",
    "        plt.imshow(image, cmap='gray')\n",
    "    else:\n",
    "        plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))\n",
    "    plt.title(title)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "# Main Processing Function\n",
    "def process_and_display_plate(image_path, padding=20):\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",
    "\n",
    "    # YOLO Detection\n",
    "    results = yolo_model.predict(image)\n",
    "    plate_found = False\n",
    "    best_confidence = -1\n",
    "    best_plate_image = None\n",
    "    best_plate_box = None\n",
    "\n",
    "    for result in results:\n",
    "        boxes = result.boxes\n",
    "        if boxes is None or len(boxes) == 0:\n",
    "            continue\n",
    "        for box in boxes:\n",
    "            confidence = box.conf.item()\n",
    "            if confidence > best_confidence:\n",
    "                best_confidence = confidence\n",
    "                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())\n",
    "                x1_padded = max(x1 - padding, 0)\n",
    "                y1_padded = max(y1 - padding, 0)\n",
    "                x2_padded = min(x2 + padding, image.shape[1])\n",
    "                y2_padded = min(y2 + padding, image.shape[0])\n",
    "                plate_crop = image[y1_padded:y2_padded, x1_padded:x2_padded]\n",
    "                best_plate_image = plate_crop\n",
    "                best_plate_box = np.array([[x1_padded, y1_padded], [x2_padded, y1_padded],\n",
    "                                           [x2_padded, y2_padded], [x1_padded, y2_padded]], dtype=\"float32\")\n",
    "                plate_found = True\n",
    "\n",
    "    if plate_found:\n",
    "        # Show the detected plate\n",
    "        cv2.polylines(output_image, [np.int0(best_plate_box)], isClosed=True, color=(0, 255, 0), thickness=2)\n",
    "        show_debug_image(\"Detected Plate\", best_plate_image)\n",
    "\n",
    "        # If the plate is tilted, use minAreaRect to get the correct bounding box\n",
    "        gray = cv2.cvtColor(best_plate_image, cv2.COLOR_BGR2GRAY)\n",
    "        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n",
    "        contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "        if contours:\n",
    "            largest_contour = max(contours, key=cv2.contourArea)\n",
    "            rect = cv2.minAreaRect(largest_contour)\n",
    "            box = cv2.boxPoints(rect)\n",
    "            box = np.int0(box)\n",
    "\n",
    "            # Adjust box coordinates relative to the original image\n",
    "            box[:, 0] += int(best_plate_box[0][0])\n",
    "            box[:, 1] += int(best_plate_box[0][1])\n",
    "\n",
    "            # Apply perspective transform\n",
    "            plate_image = four_point_transform(image, box)\n",
    "\n",
    "            # Show the transformed plate\n",
    "            show_debug_image(\"Transformed Plate\", plate_image)\n",
    "\n",
    "            # Recognize the plate\n",
    "            plate_number = recognize_plate(plate_image)\n",
    "\n",
    "            # Draw the box on the output image\n",
    "            cv2.polylines(output_image, [box], isClosed=True, color=(0, 0, 255), thickness=2)\n",
    "        else:\n",
    "            print(\"No contours found in the plate area.\")\n",
    "            plate_number = None\n",
    "            plate_image = None\n",
    "    else:\n",
    "        print(\"No plates detected by YOLO.\")\n",
    "        plate_number = None\n",
    "        plate_image = None\n",
    "\n",
    "    # Display the results\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",
    "        # Display recognized plate number\n",
    "        if plate_number:\n",
    "            placeholder_text = \"Number plate:\"\n",
    "            decoded_text = plate_number\n",
    "        else:\n",
    "            placeholder_text = \"Not found\"\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'Processing Result: {image_path}')\n",
    "    plt.show()\n",
    "\n",
    "# Process Images\n",
    "images = ['img/1.jpg', 'img/2.jpg', 'img/3.jpg', 'img/ru4018185.jpg']\n",
    "\n",
    "for img_path in images:\n",
    "    process_and_display_plate(img_path, padding=20)\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
}
