viz_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ankur/anaconda3/envs/memnet/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "from src.data.dataset import Spatial_Coordination_Dataset\n",
    "from src.data.dataset import Spatial_Free_Recall_Dataset\n",
    "from src.data.dataset import Spatial_Integration_Dataset\n",
    "from src.data.dataset import Spatial_Memory_Updating_Dataset\n",
    "from src.data.dataset import Spatial_Task_Switching_Dataset\n",
    "from src.data.dataset import Complex_WM_Dataset\n",
    "from src.data.dataset import Visual_Item_Recognition_RI_2C_Dataset\n",
    "from src.data.dataset import Visual_Serial_Recall_Recognition_Dataset\n",
    "from src.data.dataset import Color_Orientation_Size_Gap_Conjunction_Change_Detection_Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Coordination Task (SC Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Spatial_Coordination Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "sc_task_dataset = Spatial_Coordination_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               show_gt_pattern=True, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 1\n",
    "\n",
    "# Load the data\n",
    "sc_task_trial = sc_task_dataset[trial_num]\n",
    "stim_seq, gt, gt_pattern_img, seq_len, symmetry_offset = sc_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as sc_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"sc_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./sc_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "Image(url='./sc_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Free Recall Task (SFR Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Spatial_Free_Recall Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "sfr_task_dataset = Spatial_Free_Recall_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "sfr_task_trial = sfr_task_dataset[trial_num]\n",
    "stim_seq, gt_one_hot, seq_len, recall_gt_orig = sfr_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as sfr_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"sfr_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./sfr_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./sfr_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Integration Task (SI Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Spatial_Integration Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "si_task_dataset = Spatial_Integration_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "si_task_trial = si_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, part_size = si_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as si_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"si_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./si_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./si_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Memory Updating Task (SMU Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Spatial_Memory_Updating Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "smu_task_dataset = Spatial_Memory_Updating_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "smu_task_trial = smu_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, set_size = smu_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as smu_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"smu_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./smu_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./smu_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual Serial Recall Task (VSR Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Visual_Serial_Recall Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "vsr_task_dataset = Visual_Serial_Recall_Recognition_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               probe_variant='Recall', \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "vsr_task_trial = vsr_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, list_length = vsr_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as vsr_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"vsr_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./vsr_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./vsr_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual Serial Recognition Task (VSRec Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Visual_Serial_Recognition Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "vsrec_task_dataset = Visual_Serial_Recall_Recognition_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               probe_variant='Recognition', \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "vsrec_task_trial = vsrec_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, list_length, distractor_diff = vsrec_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as vsrec_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"vsrec_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./vsrec_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./vsrec_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual Item Recognition Task (VIR Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Visual_Item_Recognition_RI_2C Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "vir_task_dataset = Visual_Item_Recognition_RI_2C_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "vir_task_trial = vir_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, ri, gt_index = vir_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as vir_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"vir_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./vir_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./vir_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complex Span (CS Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Complex_WM Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "cs_task_dataset = Complex_WM_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "cs_task_trial = cs_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, num_distractor, variation = cs_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as cs_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"cs_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./cs_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./cs_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spatial Task Switching (STS Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Spatial_Task_Switching Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "sts_task_dataset = Spatial_Task_Switching_Dataset(data_path='./wm_bench_data/', \n",
    "                                               max_seq_len=20, \n",
    "                                               variant='Cued', \n",
    "                                               rs_img_size=96, \n",
    "                                               split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 3000\n",
    "\n",
    "# Load the data\n",
    "sts_task_trial = sts_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len = sts_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as sts_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"sts_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./sts_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./sts_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Change Detection Task (CD Task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting data for Color_Orientation_Size_Gap_Conjunction_Change_Detection_Color Task\n",
      "Data already exists. Skipping data generation.\n"
     ]
    }
   ],
   "source": [
    "cd_task_dataset = Color_Orientation_Size_Gap_Conjunction_Change_Detection_Dataset(data_path='./wm_bench_data/', \n",
    "                                                                                max_seq_len=20, \n",
    "                                                                                variant='Color',\n",
    "                                                                                rs_img_size=96, \n",
    "                                                                                split='test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "trial_num = 1600\n",
    "\n",
    "# Load the data\n",
    "cd_task_trial = cd_task_dataset[trial_num]\n",
    "stim_seq, gt, seq_len, ri, set_size = cd_task_trial\n",
    "\n",
    "# Generate gif of the stimulus sequence\n",
    "stim_seq_gif = []\n",
    "for stim in stim_seq:    \n",
    "    stim_array = stim.numpy()\n",
    "    stim_array = np.transpose(stim_array, (1, 2, 0))\n",
    "    stim_seq_gif.append(Image.fromarray((stim_array*255).astype(np.uint8)))\n",
    "\n",
    "stim_seq_gif = stim_seq_gif[:seq_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GIF generated and saved as cd_task.gif\n"
     ]
    }
   ],
   "source": [
    "gif_filename = \"cd_task.gif\"\n",
    "\n",
    "# Save the GIF\n",
    "stim_seq_gif[0].save(\n",
    "    gif_filename,\n",
    "    save_all=True,\n",
    "    append_images=stim_seq_gif[1:],\n",
    "    loop=0,\n",
    "    duration=1000,\n",
    "    disposal=2,\n",
    ")\n",
    "\n",
    "print(f\"GIF generated and saved as {gif_filename}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"./cd_task.gif\" width=\"250\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(url='./cd_task.gif', width=250)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memnet",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}