data_3mice_pp.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "from torch_geometric.data import Data, Dataset\n",
    "from torch_geometric.utils import to_networkx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the data\n",
    "\n",
    "submission_clips = np.load('data/raw/submission_data.npy',allow_pickle=True).item()\n",
    "user_train = np.load('data/raw/user_train.npy',allow_pickle=True).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# understanding the arrays\n",
    "\n",
    "sequence_names = list(user_train['sequences'].keys())\n",
    "# len(user_train['sequences'][sequence_names[1599]]['keypoints'][1799][2][11])\n",
    "# -> (1600 sequences) x (1800 keypoints) x (3 mouses) x (12 body parts) x (x, y  coordinates)\n",
    "\n",
    "# 1) nose, 2) left ear, 3) right ear, 4) neck, 5) left forepaw, 6) right forepaw, \n",
    "# 7) center back, 8) left hindpaw, 9) right hindpaw, 10) tail base, 11) tail middle, 12) tail tip.\n",
    "\n",
    "raw_features = user_train['sequences'][sequence_names[0]]['keypoints']\n",
    "features = raw_features.reshape(-1,3*12,2)\n",
    "# user_train['sequences'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# understanding the data\n",
    "\n",
    "# print(\"Dataset keys -\", submission_clips.keys())\n",
    "# print(\"Number of submission sequences -\", len(submission_clips['sequences']))\n",
    "\n",
    "# sequence_names = list(submission_clips[\"sequences\"].keys())\n",
    "# sequence_key = sequence_names[0]\n",
    "# single_sequence = submission_clips[\"sequences\"][sequence_key][\"keypoints\"]\n",
    "\n",
    "# print(\"Sequence name -\", sequence_key)\n",
    "# print(\"Single Sequence shape\", single_sequence.shape)\n",
    "# print(f\"Number of Frames in {sequence_key} -\", len(single_sequence))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# code for creating the edge indices\n",
    "\n",
    "# for the 2 mice dataset\n",
    "i = [0,0,1,1,3,3,4,5]\n",
    "j = [1,2,3,2,4,5,6,6]\n",
    "i2 = [x+7 for x in i]\n",
    "j2 = [x+7 for x in j]\n",
    "second_option = [i+[3]+i2,j+[10]+j2]\n",
    "second_option\n",
    "\n",
    "# for the 3 mice dataset\n",
    "i1 = [0,0,1,2,3,3,3,4,5,6,6,6,7,8,9,10]\n",
    "j1 = [1,2,3,3,4,5,6,6,6,7,8,9,9,9,10,11]\n",
    "i2 = [x+12 for x in i1]\n",
    "j2 = [x+12 for x in j1]\n",
    "i3 = [x+2*12 for x in i1]\n",
    "j3 = [x+2*12 for x in j1]\n",
    "\n",
    "# third_option = [i1+[3]+i2+[3+11]+i3,  # with edges in between mice\n",
    "#                 j1+[3+11]+j2+[3+2*11]+j3]\n",
    "\n",
    "third_option = [i1+i2+i3,j1+j2+j3]  # without edges in between mice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader\n",
    "\n",
    "class MABDataset(Dataset):\n",
    "    def __init__(self, root, test=False, transform=None, pre_transform=None, pre_filter=None):\n",
    "        self.test=test\n",
    "        super().__init__(root, transform, pre_transform, pre_filter)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        if self.test:\n",
    "            # raise NotImplementedError\n",
    "            return 'submission_data.npy'\n",
    "        return 'user_train.npy',\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        \"\"\" If these files are found in raw_dir, processing is skipped\"\"\"\n",
    "        return 'not_processed.pt'\n",
    "        # self.raw_data = np.load(self.raw_paths[0], allow_pickle=True)\n",
    "\n",
    "        # if self.test:\n",
    "        #     return [f'data_test_{i}.pt' for i in range(len(self.raw_data[()]['annotator-id_0']['task1/test/mouse071_task1_annotator1']['keypoints']))]\n",
    "        # else:\n",
    "        #     return [f'data_{i}.pt' for i in range(len(self.raw_data[()]['annotator-id_0']['task1/train/mouse001_task1_annotator1']['keypoints']))]\n",
    "\n",
    "    def download(self):\n",
    "        pass\n",
    "\n",
    "    def _edge_index_creator(self) -> list:\n",
    "        i1 = [0,0,1,2,3,3,3,4,5,6,6,6,7,8,9,10]\n",
    "        j1 = [1,2,3,3,4,5,6,6,6,7,8,9,9,9,10,11]\n",
    "\n",
    "        i2 = [x+12 for x in i1]\n",
    "        j2 = [x+12 for x in j1]\n",
    "        i3 = [x+2*12 for x in i1]\n",
    "        j3 = [x+2*12 for x in j1]\n",
    "\n",
    "        result = [i1+i2+i3,\n",
    "                  j1+j2+j3]\n",
    "        return result\n",
    "\n",
    "    def process(self):\n",
    "        self.raw_data = np.load(self.raw_paths[0], allow_pickle=True)\n",
    "        if self.test:\n",
    "            raise NotImplementedError\n",
    "            # raw_features = self.raw_data[()]['annotator-id_0']['task1/test/mouse071_task1_annotator1']['keypoints']\n",
    "            # self.labels = self.raw_data[()]['annotator-id_0']['task1/test/mouse071_task1_annotator1']['annotations']\n",
    "        else:\n",
    "            # print(self.raw_data[()]['sequences'].keys())\n",
    "            sequence_names = list(self.raw_data[()]['sequences'])\n",
    "            raw_features = self.raw_data[()]['sequences'][sequence_names[0]]['keypoints']\n",
    "            self.labels = [random.randrange(0,4) for _ in range(len(raw_features))]\n",
    "        self.num_clases = 4\n",
    "        features = raw_features.reshape(-1,3*12,2)\n",
    "\n",
    "        edge_index = torch.tensor(self._edge_index_creator(), dtype=torch.long)\n",
    "        for i in range(len(features)):\n",
    "            x = torch.tensor(features[i], dtype=torch.float)\n",
    "            y = torch.tensor(self.labels[i], dtype=torch.int)\n",
    "            graph = Data(x=x, edge_index=edge_index, y=y)\n",
    "            if self.test:\n",
    "                torch.save(graph, os.path.join(self.processed_dir, f'data_test_{i}.pt'))\n",
    "            else:\n",
    "                torch.save(graph, os.path.join(self.processed_dir, f'data_{i}.pt'))\n",
    "\n",
    "    def len(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def get(self, idx):\n",
    "        if self.test:\n",
    "            data = torch.load(os.path.join(self.processed_dir, f'data_test_{idx}.pt'))\n",
    "        else:\n",
    "            data = torch.load(os.path.join(self.processed_dir, f'data_{idx}.pt'))\n",
    "        return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataset\n",
    "\n",
    "train_dataset = MABDataset(root='./data')\n",
    "# test_dataset = MABDataset(root='./data',test=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the graph\n",
    "\n",
    "G = to_networkx(train_dataset[0], to_undirected=True)\n",
    "\n",
    "# method 1\n",
    "G = nx.Graph(G, length=20)\n",
    "nx.draw(G, node_size=100, alpha=0.8, arrowsize=8, with_labels=False)\n",
    "\n",
    "# method 2\n",
    "# def visualize(h, color, epoch=None, loss=None):\n",
    "#     plt.figure(figsize=(7,7))\n",
    "#     plt.xticks([])\n",
    "#     plt.yticks([])\n",
    "\n",
    "#     if torch.is_tensor(h):\n",
    "#         h = h.detach().cpu().numpy()\n",
    "#         plt.scatter(h[:, 0], h[:, 1], s=140, c=color, cmap=\"Set2\")\n",
    "#         if epoch is not None and loss is not None:\n",
    "#             plt.xlabel(f'Epoch: {epoch}, Loss: {loss.item():.4f}', fontsize=16)\n",
    "#     else:\n",
    "#         nx.draw_networkx(G, pos=nx.spring_layout(G, seed=42), with_labels=False,\n",
    "#                          node_color=color, cmap=\"Set2\")\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''code from the aicrowd webpage\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "from matplotlib import colors\n",
    "from matplotlib import rc\n",
    "\n",
    "rc('animation', html='jshtml')\n",
    "\n",
    "# Note: Image processing may be slow if too many frames are animated.\n",
    "\n",
    "#Plotting constants\n",
    "FRAME_WIDTH_TOP = 850\n",
    "FRAME_HEIGHT_TOP = 850\n",
    "\n",
    "M1_COLOR = 'lawngreen'\n",
    "M2_COLOR = 'skyblue'\n",
    "M3_COLOR = 'tomato'\n",
    "\n",
    "PLOT_MOUSE_START_END = [(0, 1), (1, 3), (3, 2), (2, 0),        # head\n",
    "                        (3, 6), (6, 9),                        # midline\n",
    "                        (9, 10), (10, 11),                     # tail\n",
    "                        (4, 5), (5, 8), (8, 9), (9, 7), (7, 4) # legs\n",
    "                       ]\n",
    "\n",
    "class_to_number = {s: i for i, s in enumerate(user_train['vocabulary'])}\n",
    "\n",
    "number_to_class = {i: s for i, s in enumerate(user_train['vocabulary'])}\n",
    "\n",
    "def num_to_text(anno_list):\n",
    "  return np.vectorize(number_to_class.get)(anno_list)\n",
    "\n",
    "def set_figax():\n",
    "    fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "    img = np.zeros((FRAME_HEIGHT_TOP, FRAME_WIDTH_TOP, 3))\n",
    "\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.imshow(img)\n",
    "\n",
    "    ax.get_xaxis().set_visible(False)\n",
    "    ax.get_yaxis().set_visible(False)\n",
    "\n",
    "    return fig, ax\n",
    "\n",
    "def plot_mouse(ax, pose, color):\n",
    "    # Draw each keypoint\n",
    "    for j in range(10):\n",
    "        ax.plot(pose[j, 0], pose[j, 1], 'o', color=color, markersize=3)\n",
    "\n",
    "    # Draw a line for each point pair to form the shape of the mouse\n",
    "\n",
    "    for pair in PLOT_MOUSE_START_END:\n",
    "        line_to_plot = pose[pair, :]\n",
    "        ax.plot(line_to_plot[:, 0], line_to_plot[\n",
    "                :, 1], color=color, linewidth=1)\n",
    "\n",
    "def animate_pose_sequence(video_name, seq, start_frame = 0, stop_frame = 100, skip = 0,\n",
    "                          annotation_sequence = None):\n",
    "    # Returns the animation of the keypoint sequence between start frame\n",
    "    # and stop frame. Optionally can display annotations.\n",
    "\n",
    "    image_list = []\n",
    "\n",
    "    counter = 0\n",
    "    if skip:\n",
    "        anim_range = range(start_frame, stop_frame, skip)\n",
    "    else:\n",
    "        anim_range = range(start_frame, stop_frame)\n",
    "\n",
    "    for j in anim_range:\n",
    "        if counter%20 == 0:\n",
    "          print(\"Processing frame \", j)\n",
    "        fig, ax = set_figax()\n",
    "        plot_mouse(ax, seq[j, 0, :, :], color=M1_COLOR)\n",
    "        plot_mouse(ax, seq[j, 1, :, :], color=M2_COLOR)\n",
    "        plot_mouse(ax, seq[j, 2, :, :], color=M3_COLOR)\n",
    "\n",
    "        if annotation_sequence is not None:\n",
    "          annot = annotation_sequence[j]\n",
    "          annot = number_to_class[annot]\n",
    "          plt.text(50, -20, annot, fontsize = 16,\n",
    "                   bbox=dict(facecolor=class_to_color[annot], alpha=0.5))\n",
    "\n",
    "        ax.set_title(\n",
    "            video_name + '\\n frame {:03d}.png'.format(j))\n",
    "\n",
    "        ax.axis('off')\n",
    "        fig.tight_layout(pad=0)\n",
    "        ax.margins(0)\n",
    "\n",
    "        fig.canvas.draw()\n",
    "        image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(),\n",
    "                                        dtype=np.uint8)\n",
    "        image_from_plot = image_from_plot.reshape(\n",
    "            fig.canvas.get_width_height()[::-1] + (3,))\n",
    "\n",
    "        image_list.append(image_from_plot)\n",
    "\n",
    "        plt.close()\n",
    "        counter = counter + 1\n",
    "\n",
    "    # Plot animation.\n",
    "    fig = plt.figure(figsize=(8,8))\n",
    "    plt.axis('off')\n",
    "    im = plt.imshow(image_list[0])\n",
    "\n",
    "    def animate(k):\n",
    "        im.set_array(image_list[k])\n",
    "        return im,\n",
    "    ani = animation.FuncAnimation(fig, animate, frames=len(image_list), blit=True)\n",
    "    return ani\n",
    "    '''"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8b65681a3d19f4b473b45a90bc0c881aec513e0960c214c85a14e16ba4a66a1e"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('cs159')",
   "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.9.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}