sample_submission.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "about-heavy",
   "metadata": {},
   "source": [
    "## 0. Libarary 불러오기 및 경로설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cubic-scoop",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from torchvision import transforms\n",
    "from torchvision.transforms import Resize, ToTensor, Normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "built-elevation",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 테스트 데이터셋 폴더 경로를 지정해주세요.\n",
    "test_dir = '/mnt/ssd/data/mask_final/public'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "quiet-organizer",
   "metadata": {},
   "source": [
    "## 1. Model 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "acknowledged-easter",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(nn.Module):\n",
    "    def __init__(self, num_classes: int = 1000):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Dropout(),\n",
    "            nn.Linear(64, 32),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(32, num_classes),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.features(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "domestic-channels",
   "metadata": {},
   "source": [
    "## 2. Test Dataset 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "extensive-north",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestDataset(Dataset):\n",
    "    def __init__(self, img_paths, transform):\n",
    "        self.img_paths = img_paths\n",
    "        self.transform = transform\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image = Image.open(self.img_paths[index])\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "continued-feelings",
   "metadata": {},
   "source": [
    "## 3. Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "coral-shade",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test inference is done!\n"
     ]
    }
   ],
   "source": [
    "# meta 데이터와 이미지 경로를 불러옵니다.\n",
    "submission = pd.read_csv(os.path.join(test_dir, 'info.csv'))\n",
    "image_dir = os.path.join(test_dir, 'images')\n",
    "\n",
    "# Test Dataset 클래스 객체를 생성하고 DataLoader를 만듭니다.\n",
    "image_paths = [os.path.join(image_dir, img_id) for img_id in submission.ImageID]\n",
    "transform = transforms.Compose([\n",
    "    Resize((512, 384), Image.BILINEAR),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=(0.5, 0.5, 0.5), std=(0.2, 0.2, 0.2)),\n",
    "])\n",
    "dataset = TestDataset(image_paths, transform)\n",
    "\n",
    "loader = DataLoader(\n",
    "    dataset,\n",
    "    shuffle=False\n",
    ")\n",
    "\n",
    "# 모델을 정의합니다. (학습한 모델이 있다면 torch.load로 모델을 불러주세요!)\n",
    "device = torch.device('cuda')\n",
    "model = MyModel(num_classes=18).to(device)\n",
    "model.eval()\n",
    "\n",
    "# 모델이 테스트 데이터셋을 예측하고 결과를 저장합니다.\n",
    "all_predictions = []\n",
    "for images in loader:\n",
    "    with torch.no_grad():\n",
    "        images = images.to(device)\n",
    "        pred = model(images)\n",
    "        pred = pred.argmax(dim=-1)\n",
    "        all_predictions.extend(pred.cpu().numpy())\n",
    "submission['ans'] = all_predictions\n",
    "\n",
    "# 제출할 파일을 저장합니다.\n",
    "submission.to_csv(os.path.join(test_dir, 'submission.csv'), index=False)\n",
    "print('test inference is done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "verbal-sample",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}