{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a8ca6784",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fri Mar 27 08:05:49 2026       \n",
      "+-----------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 580.82.07              Driver Version: 580.82.07      CUDA Version: 13.0     |\n",
      "+-----------------------------------------+------------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                        |               MIG M. |\n",
      "|=========================================+========================+======================|\n",
      "|   0  Tesla T4                       Off |   00000000:00:04.0 Off |                    0 |\n",
      "| N/A   67C    P0             31W /   70W |     223MiB /  15360MiB |      0%      Default |\n",
      "|                                         |                        |                  N/A |\n",
      "+-----------------------------------------+------------------------+----------------------+\n",
      "\n",
      "+-----------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                              |\n",
      "|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |\n",
      "|        ID   ID                                                               Usage      |\n",
      "|=========================================================================================|\n",
      "|    0   N/A  N/A            2100      C   /usr/bin/python3                        220MiB |\n",
      "+-----------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0bb11592",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms, datasets, models\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "import os\n",
    "import shutil\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "054b3dad",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c78b4c4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.listdir('/content/drive/MyDrive/TrashNet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "598ad20a",
   "metadata": {},
   "outputs": [],
   "source": [
    "source = '/content/drive/MyDrive/TrashNet'\n",
    "base = '/content/drive/MyDrive/waste_dataset'\n",
    "\n",
    "train_dir = os.path.join(base, 'train')\n",
    "val_dir = os.path.join(base, 'val')\n",
    "\n",
    "os.makedirs(train_dir, exist_ok=True)\n",
    "os.makedirs(val_dir, exist_ok=True)\n",
    "\n",
    "for class_name in os.listdir(source):\n",
    "    class_path = os.path.join(source, class_name)\n",
    "    \n",
    "    images = os.listdir(class_path)\n",
    "    random.shuffle(images)\n",
    "\n",
    "    split = int(0.8 * len(images))\n",
    "\n",
    "    train_images = images[:split]\n",
    "    val_images = images[split:]\n",
    "\n",
    "    os.makedirs(os.path.join(train_dir, class_name), exist_ok=True)\n",
    "    os.makedirs(os.path.join(val_dir, class_name), exist_ok=True)\n",
    "\n",
    "    for img in train_images:\n",
    "        shutil.copy(os.path.join(class_path, img),\n",
    "                    os.path.join(train_dir, class_name, img))\n",
    "\n",
    "    for img in val_images:\n",
    "        shutil.copy(os.path.join(class_path, img),\n",
    "                    os.path.join(val_dir, class_name, img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a743b4fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']\n",
      "2422\n",
      "911\n"
     ]
    }
   ],
   "source": [
    "train_path = '/content/drive/MyDrive/waste_dataset/train'\n",
    "val_path = '/content/drive/MyDrive/waste_dataset/val'\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.Resize((224,224)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomRotation(15),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485,0.456,0.406],\n",
    "                         [0.229,0.224,0.225])\n",
    "])\n",
    "\n",
    "val_transform = transforms.Compose([\n",
    "    transforms.Resize((224,224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485,0.456,0.406],\n",
    "                         [0.229,0.224,0.225])\n",
    "])\n",
    "\n",
    "from torchvision import datasets\n",
    "\n",
    "train_data = datasets.ImageFolder(train_path, transform=train_transform)\n",
    "val_data = datasets.ImageFolder(val_path, transform=val_transform)\n",
    "\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_loader = DataLoader(train_data, batch_size=32, shuffle=True)\n",
    "val_loader = DataLoader(val_data, batch_size=32)\n",
    "\n",
    "\n",
    "print(train_data.classes)\n",
    "print(len(train_data))\n",
    "print(len(val_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "01fc969b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linear(in_features=2048, out_features=6, bias=True)\n"
     ]
    }
   ],
   "source": [
    "from torchvision import models\n",
    "import torch.nn as nn\n",
    "\n",
    "model = models.resnet50(pretrained=True)\n",
    "\n",
    "model.fc = nn.Linear(model.fc.in_features, 6)\n",
    "\n",
    "import torch.optim as optim\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "print(model.fc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "17fede0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Freeze\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# Replace final layer\n",
    "model.fc = nn.Linear(model.fc.in_features, 6)\n",
    "\n",
    "# Move to GPU\n",
    "model = model.to(device)\n",
    "\n",
    "# Optimizer (IMPORTANT)\n",
    "optimizer = optim.Adam(model.fc.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b441ea41",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/3], Loss: 1.1354\n",
      "Validation Accuracy: 69.48%\n",
      "Epoch [2/3], Loss: 0.7224\n",
      "Validation Accuracy: 79.80%\n",
      "Epoch [3/3], Loss: 0.6208\n",
      "Validation Accuracy: 82.44%\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 3\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "\n",
    "    for images, labels in train_loader:\n",
    "        images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "    print(f\"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}\")\n",
    "\n",
    "    \n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in val_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    print(f\"Validation Accuracy: {100 * correct / total:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1a5ce395",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model Saved Succesfully\n",
      "['.config', 'waste_classifier.pth', 'drive', 'sample_data']\n"
     ]
    }
   ],
   "source": [
    "torch.save(model.state_dict(), \"waste_classifier.pth\")\n",
    "print(\"Model Saved Succesfully\")\n",
    "\n",
    "print(os.listdir())"
   ]
  }
 ],
 "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}