notebooks/shape_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Train A Shape Classifier Model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "train_data_root = \"../datasets/train\"\n",
    "test_data_root = \"../datasets/test\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import os\n",
    "\n",
    "# Set device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Define transformations (including resizing and normalization)\n",
    "transform = transforms.Compose([\n",
    "    transforms.Grayscale(num_output_channels=1),  # Convert to grayscale (black and white images)\n",
    "    transforms.Resize((64, 64)),  # Resize images to 64x64 pixels\n",
    "    transforms.ToTensor(),  # Convert the image to a tensor\n",
    "    transforms.Normalize((0.5,), (0.5,))  # Normalize the images (mean=0.5, std=0.5 for grayscale)\n",
    "])\n",
    "\n",
    "# Load the dataset\n",
    "train_dataset = datasets.ImageFolder(root=train_data_root, transform=transform)\n",
    "test_dataset = datasets.ImageFolder(root=test_data_root, transform=transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
    "\n",
    "# Check class names (optional)\n",
    "print(f'Classes: {train_dataset.classes}')\n",
    "\n",
    "# 2. Define a simple CNN model\n",
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=1)\n",
    "        self.fc1 = nn.Linear(32 * 16 * 16, 128)\n",
    "        self.fc2 = nn.Linear(128, 3)  # 3 classes: circle, triangle, rectangle\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))   # First Conv Layer\n",
    "        x = F.max_pool2d(x, 2)      # Max Pooling\n",
    "        x = F.relu(self.conv2(x))   # Second Conv Layer\n",
    "        x = F.max_pool2d(x, 2)      # Max Pooling\n",
    "        x = x.view(x.size(0), -1)   # Flatten\n",
    "        x = F.relu(self.fc1(x))     # Fully Connected Layer 1\n",
    "        x = self.fc2(x)             # Fully Connected Layer 2 (output)\n",
    "        return x\n",
    "\n",
    "model = SimpleCNN().to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "def train_model(model, train_loader, criterion, optimizer, epochs=10):\n",
    "    model.train()\n",
    "    for epoch in range(epochs):\n",
    "        running_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            \n",
    "            # Zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # Forward pass\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            # Backward pass and optimize\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            # Statistics\n",
    "            running_loss += loss.item()\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "        \n",
    "        accuracy = 0 * correct / total\n",
    "        print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}, Accuracy: {accuracy:.2f}%')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_model(model, train_loader, criterion, optimizer, epochs=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "# Test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, test_loader):\n",
    "    \"\"\"Print the Precision, Recall and F1-score for the trained model\n",
    "    \"\"\"\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "# Show Predictions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_prediction(model, image):\n",
    "    \"\"\"Pass the image to the model and overlay the predicted shape and confidence on the input\n",
    "    image and display it\n",
    "    \"\"\"\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}