old_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attempts to create a pose estimation for spacecraft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive approach using existing basic Linear model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a dataset per spacecraft."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of spacecraft: 660\n"
     ]
    }
   ],
   "source": [
    "# create dataset from train_labels.csv\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import os\n",
    "import io\n",
    "from skimage import io, transform\n",
    "import random\n",
    "# create dataset that contains the image label and data label\n",
    "class SpacecraftPoseDataset(Dataset):\n",
    "    \"\"\" Spacecraft pose dataset \"\"\"\n",
    "\n",
    "    def __init__(self, csv_file=None, root_dir=None, df=None, transform = None):\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            csv_file (string): Path to the csv file with annotations.\n",
    "            root_dir (string): Directory with all the images.\n",
    "            transform (callable, optional): Optional transform to be applied\n",
    "                on a sample.\n",
    "        \"\"\"\n",
    "        if df is not None:\n",
    "            self.labels_frame = df\n",
    "        else:\n",
    "            self.labels_frame = pd.read_csv(csv_file)\n",
    "        self.root_dir = root_dir\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels_frame)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "\n",
    "        img_name = os.path.join(self.root_dir, self.labels_frame.iloc[idx, 0], f'{self.labels_frame.iloc[idx, 1]:03}.png')\n",
    "        image = io.imread(img_name)\n",
    "        label = self.labels_frame.iloc[idx, 1:]\n",
    "        sample = {'image': image, 'label': label}\n",
    "\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "\n",
    "        return sample\n",
    "\n",
    "df = pd.read_csv('data/train_labels.csv')\n",
    "# split df by chain_id\n",
    "spacecraft = []\n",
    "# print(df.head())\n",
    "for chain_id in df['chain_id'].unique():\n",
    "    spacecraft.append(df[df['chain_id'] == chain_id])\n",
    "\n",
    "# print(spacecraft)\n",
    "for i in range(len(spacecraft)):\n",
    "    spacecraft[i] = SpacecraftPoseDataset(df=spacecraft[i], root_dir='data/images')\n",
    "    # print(spacecraft[i].__len__())\n",
    "    # print(spacecraft[i].__getitem__(0))\n",
    "\n",
    "# SpacecraftPoseDataset = SpacecraftPoseDataset(csv_file = 'data/train_labels.csv', root_dir = 'data/images')\n",
    "\n",
    "# train_dataloader = DataLoader(SpacecraftPoseDataset, batch_size=64)\n",
    "\n",
    "# split dataset\n",
    "# train_set, test_set, unused_set = torch.utils.data.random_split(SpacecraftPoseDataset, [0.2, 0.05, 0.75])\n",
    "print(\"num of spacecraft:\",len(spacecraft))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use torch to make a Linear model\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "class LinearModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LinearModel, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(1024, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 7),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # resize x to 3x1280x1024\n",
    "        # x = x.reshape(3, 1280, 1024)\n",
    "       \n",
    "        # convert x to tensor\n",
    "        x = torch.tensor(x, dtype=torch.float32, device='mps')\n",
    "        # take the average of the 3 channels\n",
    "        x = torch.mean(x, dim=-1)\n",
    "        x = torch.mean(x, dim=-1)\n",
    "\n",
    "        # print(x.shape)\n",
    "        x = self.model(x)\n",
    "        return x\n",
    "\n",
    "# Create an instance of the custom CNN model\n",
    "cnn_model = LinearModel()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use torch to make a CNN model\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "class BasicCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BasicCNN, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),\n",
    "            nn.MaxPool2d(kernel_size=10, stride=5),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=10, stride=5),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.model2 = nn.Sequential(\n",
    "            nn.Linear(29184, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 7),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # resize x to 3x1280x1024\n",
    "        x = x.reshape(3, 1280, 1024)\n",
    "        # convert x to tensor\n",
    "        x = torch.tensor(x, dtype=torch.float32, device='mps')\n",
    "        # print(x.shape)\n",
    "        x = self.model(x)\n",
    "        # x = x.reshape(7)\n",
    "        x = self.flatten(x)\n",
    "        x = x.reshape(29184)\n",
    "        x = self.model2(x)  \n",
    "        return x\n",
    "\n",
    "# Create an instance of the custom CNN model\n",
    "cnn_model = BasicCNN()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use torch to make a simpler CNN model\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "class BasicCNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BasicCNN, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=20, stride=10),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=20, stride=10),\n",
    "        )\n",
    "        self.model2 = nn.Sequential(\n",
    "            nn.Linear(3168, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 7),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # resize x to 3x1280x1024\n",
    "        x = x.reshape(3, 1280, 1024)\n",
    "        # convert x to tensor\n",
    "        x = torch.tensor(x, dtype=torch.float32, device='mps')\n",
    "        # print(x.shape)\n",
    "        x = self.model(x)\n",
    "        # x = x.reshape(7)\n",
    "        x = self.flatten(x)\n",
    "        x = x.reshape(3168)\n",
    "        x = self.model2(x)  \n",
    "        return x\n",
    "\n",
    "# Create an instance of the custom CNN model\n",
    "cnn_model = BasicCNN()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "Time taken so far: 0.0001780986785888672\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 2/10\n",
      "Time taken so far: 152.5457682609558\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 3/10\n",
      "Time taken so far: 300.20884704589844\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 4/10\n",
      "Time taken so far: 446.38360023498535\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 5/10\n",
      "Time taken so far: 595.7526142597198\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 6/10\n",
      "Time taken so far: 741.5646421909332\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 7/10\n",
      "Time taken so far: 889.7017440795898\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 8/10\n",
      "Time taken so far: 1037.802971124649\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 9/10\n",
      "Time taken so far: 1183.5211880207062\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n",
      "Epoch 10/10\n",
      "Time taken so far: 1337.284977197647\n",
      "Batch 1/30\n",
      "Batch 2/30\n",
      "Batch 3/30\n",
      "Batch 4/30\n",
      "Batch 5/30\n",
      "Batch 6/30\n",
      "Batch 7/30\n",
      "Batch 8/30\n",
      "Batch 9/30\n",
      "Batch 10/30\n",
      "Batch 11/30\n",
      "Batch 12/30\n",
      "Batch 13/30\n",
      "Batch 14/30\n",
      "Batch 15/30\n",
      "Batch 16/30\n",
      "Batch 17/30\n",
      "Batch 18/30\n",
      "Batch 19/30\n",
      "Batch 20/30\n",
      "Batch 21/30\n",
      "Batch 22/30\n",
      "Batch 23/30\n",
      "Batch 24/30\n",
      "Batch 25/30\n",
      "Batch 26/30\n",
      "Batch 27/30\n",
      "Batch 28/30\n",
      "Batch 29/30\n",
      "Batch 30/30\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "# dataset = SpacecraftPoseDataset\n",
    "dataset = spacecraft\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(cnn_model.parameters(), lr=0.0001)\n",
    "\n",
    "# Set the model to training mode\n",
    "cnn_model.train()\n",
    "\n",
    "# Train the model\n",
    "# Set the number of epochs\n",
    "num_epochs = 10\n",
    "\n",
    "num_spacecrafts = 30\n",
    "\n",
    "# Set the initial running loss\n",
    "running_loss = 0.0\n",
    "\n",
    "# Train the model for the specified number of epochs\n",
    "# dont use dataloader, use dataset\n",
    "\n",
    "# use m1 pro gpu to train\n",
    "mps_device = torch.device('mps')\n",
    "cnn_model.to(mps_device)\n",
    "startTime = time.time()\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    print(f'Epoch {epoch+1}/{num_epochs}')\n",
    "    print(\"Time taken so far:\", time.time() - startTime)\n",
    "    for i in range(num_spacecrafts):\n",
    "        print(f'Batch {i+1}/{num_spacecrafts}')\n",
    "        i = random.randint(0, len(dataset) - 1)\n",
    "        data = dataset[i]\n",
    "\n",
    "        for j in range(len(data)):\n",
    "            sample = data[j]\n",
    "            image = sample['image']\n",
    "            label = sample['label']\n",
    "\n",
    "            # convert label to tensor\n",
    "            label = torch.tensor(label[1:], dtype=torch.float32, device='mps')\n",
    "            # print(label)\n",
    "            # Zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = cnn_model(image)\n",
    "            # print(outputs.shape)\n",
    "\n",
    "            # Calculate the loss\n",
    "            loss = criterion(outputs, label)\n",
    "\n",
    "            # Backward pass\n",
    "            loss.backward()\n",
    "\n",
    "            # Optimize the weights\n",
    "            optimizer.step()\n",
    "\n",
    "            # Print the statistics\n",
    "            running_loss += loss.item()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the model\n",
    "torch.save(cnn_model.state_dict(), 'linear_model.pth')\n",
    "\n",
    "# convert to cpu\n",
    "cnn_model.to('cpu')\n",
    "# save the model\n",
    "torch.save(cnn_model.state_dict(), 'linear_model_cpu.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  chain_id  i            x          y            z           qw            qx  \\\n",
      "0        0  0  3392.112793 -53.273254 -4633.512695  2497.546143 -11770.953125   \n",
      "1        1  1  3708.927979 -57.926975 -5068.404297  2731.128174 -12872.385742   \n",
      "2        2  2  3506.833252 -55.030956 -4790.500000  2582.021729 -12169.142578   \n",
      "3        3  3  4187.043457 -65.468628 -5722.483398  3083.154297 -14532.419922   \n",
      "4        4  4  4553.538086 -71.295242 -6223.434570  3352.858154 -15804.304688   \n",
      "\n",
      "            qy            qz  \n",
      "0 -4489.433594  -9681.136719  \n",
      "1 -4909.356934 -10587.912109  \n",
      "2 -4641.238770 -10008.919922  \n",
      "3 -5542.458008 -11953.802734  \n",
      "4 -6027.540527 -13000.271484  \n"
     ]
    }
   ],
   "source": [
    "# test the model\n",
    "# Set the model to evaluation mode\n",
    "cnn_model.eval()\n",
    "\n",
    "cnn_model.to(mps_device)\n",
    "\n",
    "# Test the model\n",
    "# Set the initial running loss\n",
    "running_loss = 0.0\n",
    "\n",
    "# run the model on the 5 random spacecrafts\n",
    "for i in range(5):\n",
    "    i = random.randint(0, len(dataset) - 1)\n",
    "    sample = dataset[i]\n",
    "    result_df = pd.DataFrame(columns=['chain_id','i',\"x\", \"y\", \"z\", \"qw\", \"qx\", \"qy\", \"qz\"])\n",
    "    for j in range(len(sample)):\n",
    "        sample = data[j]\n",
    "        image = sample['image']\n",
    "        label = sample['label']\n",
    "\n",
    "        # convert label to tensor\n",
    "        label = torch.tensor(label[1:], dtype=torch.float32, device=mps_device)\n",
    "        # print(label)\n",
    "        # Forward pass\n",
    "        outputs = cnn_model(image)\n",
    "        # print(outputs.detach().numpy())\n",
    "        # add the result to the dataframe\n",
    "        temp = outputs.cpu().detach().numpy()\n",
    "        result_df = pd.concat([result_df, pd.DataFrame({'chain_id': sample['label'][0], 'i': j, \n",
    "                    'x': temp[0], 'y': temp[1], 'z': temp[2], 'qw': temp[3], \n",
    "                    'qx': temp[4], 'qy': temp[5], 'qz': temp[6]}, index=[0])], ignore_index=True)\n",
    "\n",
    "\n",
    "        # Calculate the loss\n",
    "        loss = criterion(outputs, label)\n",
    "\n",
    "        # Print the statistics\n",
    "        running_loss += loss.item()\n",
    "\n",
    "print(result_df.head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}