opencv.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wJgxVHkxMfjS"
      },
      "outputs": [],
      "source": [
        "#条件：\n",
        "#下記のコードで学習してからではないと関数predictが動かないので、一度実行をお願いします。\n",
        "#学習時間が長いのでもし短縮したい場合、MAX_EPOCH = 5の数字を小さくしてください。\n",
        "#関数predictの引数は画像のファイル名が入るります。\n",
        "#この関数は画像に写っているものの推察をしてくれますが、飛行機、車、鳥、猫、鹿、犬、カエル、馬、船、トラックのいずれかしか判断できません。\n",
        "import torch\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Conv2D\n",
        "from keras.layers import Activation, Flatten, Dense\n",
        "from keras.optimizers import Adam\n",
        "\n",
        "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])\n",
        "\n",
        "# 学習データ\n",
        "train_data_with_teacher_labels = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
        "train_data_loader = torch.utils.data.DataLoader(train_data_with_teacher_labels, batch_size=4, shuffle=True, num_workers=2)\n",
        "\n",
        "# 検証データ\n",
        "test_data_with_teacher_labels = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)\n",
        "test_data_loader = torch.utils.data.DataLoader(test_data_with_teacher_labels, batch_size=4, shuffle=False, num_workers=2)\n",
        "\n",
        "class CNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
        "        self.layer1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.layer2 = nn.Linear(120, 84)\n",
        "        self.layer3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, input_data):\n",
        "        input_data = self.pool(F.relu(self.conv1(input_data)))\n",
        "        input_data = self.pool(F.relu(self.conv2(input_data)))\n",
        "        input_data = input_data.view(-1, 16 * 5 * 5)\n",
        "        input_data = F.relu(self.layer1(input_data))\n",
        "        input_data = F.relu(self.layer2(input_data))\n",
        "        input_data = self.layer3(input_data)\n",
        "        return input_data\n",
        "\n",
        "\n",
        "model = Sequential()\n",
        "\n",
        "model.add(Conv2D(filters=10, kernel_size=(3, 3), padding='same', input_shape=(32,32,3), activation='relu'))\n",
        "model.add(Flatten())\n",
        "model.add(Dense(10, activation='softmax'))\n",
        "\n",
        "model.compile(loss='categorical_crossentropy', optimizer=Adam(), metrics=['accuracy'])\n",
        "\n",
        "# エポック数\n",
        "MAX_EPOCH = 5\n",
        "\n",
        "\n",
        "# ニューラルネットワークの定義\n",
        "class CNN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CNN, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
        "        self.pool = nn.MaxPool2d(2, 2)\n",
        "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
        "        self.layer1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.layer2 = nn.Linear(120, 84)\n",
        "        self.layer3 = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, input_data):\n",
        "        input_data = self.pool(F.relu(self.conv1(input_data)))\n",
        "        input_data = self.pool(F.relu(self.conv2(input_data)))\n",
        "        input_data = input_data.view(-1, 16 * 5 * 5)\n",
        "        input_data = F.relu(self.layer1(input_data))\n",
        "        input_data = F.relu(self.layer2(input_data))\n",
        "        input_data = self.layer3(input_data)\n",
        "        return input_data\n",
        "\n",
        "# transform定義\n",
        "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
        "\n",
        "# 学習データ\n",
        "train_data_with_teacher_labels = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
        "train_data_loader = torch.utils.data.DataLoader(train_data_with_teacher_labels, batch_size=4, shuffle=True, num_workers=2)\n",
        "\n",
        "model = CNN()\n",
        "# optimizerの設定\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.SGD(params=model.parameters(), lr=0.001, momentum=0.9)\n",
        "\n",
        "# 学習\n",
        "for epoch in range(MAX_EPOCH):\n",
        "    total_loss = 0.0\n",
        "    for i, data in enumerate(train_data_loader, 0):\n",
        "        # 学習データと教師ラベルデータを取得\n",
        "        train_data, teacher_labels = data\n",
        "        # 勾配情報を削除\n",
        "        optimizer.zero_grad()\n",
        "        # モデルで予測を計算\n",
        "        outputs = model(train_data)\n",
        "        # 微分計算\n",
        "        loss = criterion(outputs, teacher_labels)\n",
        "        loss.backward()\n",
        "        # 勾配を更新\n",
        "        optimizer.step()\n",
        "        # 誤差\n",
        "        total_loss += loss.item()\n",
        "        # 1000ミニバッチずつ進捗を表示\n",
        "        if i % 1000 == 999:\n",
        "            print('学習進捗：[学習回数：%d, ミニバッチ数：%5d] loss: %.3f' % (epoch + 1, i + 1, total_loss / 1000))\n",
        "            total_loss = 0.0\n",
        "\n",
        "# モデル保存\n",
        "torch.save(model, \"model.pth\")\n",
        "print(\"-----学習完了-----\")\n",
        "\n",
        "def predict(img):\n",
        "  import cv2\n",
        "  import matplotlib.pyplot as plt\n",
        "  import torch\n",
        "  import torchvision\n",
        "  import torchvision.transforms as transforms\n",
        "  import torch.nn as nn\n",
        "  import torch.nn.functional as F\n",
        "  import torch.optim as optim\n",
        "  from torchvision import models\n",
        "  from PIL import Image\n",
        "  import matplotlib.pyplot as plt\n",
        "\n",
        "  model = torch.load('model.pth')\n",
        "  classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')\n",
        "\n",
        "  # 推論モード\n",
        "  model.eval()\n",
        "\n",
        "  # 画像をモデルと形式を合わせる\n",
        "  sample = cv2.imread(img)\n",
        "  sample_re = cv2.resize(sample, dsize=(32,32))\n",
        "  plt.imshow(cv2.cvtColor(sample, cv2.COLOR_BGR2RGB))\n",
        "  sample_t = torchvision.transforms.functional.to_tensor(sample_re)\n",
        "  x = sample_t.unsqueeze(0)\n",
        "\n",
        "  # 推論開始\n",
        "  y_predict = model(x)\n",
        "  print(y_predict)\n",
        "  print(\"推論結果: \", classes[y_predict.argmax()])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#フォルダにsample.jpgを入れてから実行してください。\n",
        "predict(\"sample.jpg\")"
      ],
      "metadata": {
        "id": "OwGn6e-u7MfG"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}