Comparison_Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<contextlib.ExitStack at 0x7f1ac466d970>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from aicspylibczi import CziFile\n",
    "import czifile\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import imshow\n",
    "import cv2\n",
    "import os\n",
    "import imageio\n",
    "import ffmpeg\n",
    "import time\n",
    "import pandas as pd\n",
    "from cellpose import io, models\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import torch.backends.cudnn as cudnn\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import os\n",
    "import copy\n",
    "import json\n",
    "import glob\n",
    "import itertools\n",
    "from PIL import Image\n",
    "import torch.nn.functional as F\n",
    "from utils import *\n",
    "cudnn.benchmark = True\n",
    "plt.ion()   # interactive mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class VideoDataMIP:\n",
    "    def __init__(self, files):\n",
    "        self.data = {\n",
    "        }\n",
    "        \n",
    "        for category, num in files:\n",
    "            print(f\"Loading in MIP {num}\")\n",
    "            assert category == 'mip', \"Can't load non Mip file\"\n",
    "            file = {}\n",
    "            file['video'] = get_file(category, num)\n",
    "            \n",
    "            frames, shp = file['video'].read_image(C=0)\n",
    "            frames = scale_img(frames.squeeze())\n",
    "            file['frames'] = frames\n",
    "            print(f\"frames {num}: {frames.shape}\")\n",
    "            file['masks'] = binarize_video(frames)           \n",
    "    \n",
    "            self.data[num] = file    \n",
    "    def extract_all_traces(self, file_num, sequence_length, hist_length=2):\n",
    "        # hist length is how many frames of history\n",
    "        frames, masks = self.data[file_num]['frames'], self.data[file_num]['masks']\n",
    "        N = len(frames)\n",
    "        s = 0\n",
    "        all_traces = []\n",
    "        all_videos = []\n",
    "        for i in range(N // sequence_length):\n",
    "            print(f\"Extracting traces from {s}:{s+sequence_length}\")\n",
    "            data, videos = extract_traces(frames[s:s+sequence_length], masks[s:s+sequence_length], hist=hist_length)\n",
    "            s += sequence_length\n",
    "            all_traces = all_traces + data\n",
    "            all_videos = all_videos + videos\n",
    "        \n",
    "        if(N % sequence_length > 0):\n",
    "            data, videos = extract_traces(frames[-1*sequence_length:], masks[-1*sequence_length:], hist=hist_length)\n",
    "            all_traces = all_traces + data\n",
    "            all_videos = all_videos + videos        \n",
    "        self.data[file_num]['traces'] = all_traces\n",
    "        self.data[file_num]['trace_videos'] = all_videos\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fake_data = [\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones([100, 100])],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones([100, 100])]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     },\n",
    "#     {\n",
    "#         \"box\": [(1, 2, 3, 4), (1, 2, 3, 5)],\n",
    "#         \"masks\": [np.ones([100, 100]), np.ones((100, 100))],\n",
    "#         \"patch\": [np.ones([100, 100]), np.ones((100, 100))]\n",
    "#     }\n",
    "    \n",
    "# ]\n",
    "\n",
    "# fake_data[0][\"masks\"][1].flatten().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.measure import centroid\n",
    "import skimage.measure as skm\n",
    "\n",
    "\n",
    "\n",
    "box_shape = (180, 180) #TO DO: find the biggest box and set it to this\n",
    "X = 10\n",
    "\n",
    "class CellBoxMaskPatch(torch.utils.data.Dataset):\n",
    "    #input will be a Directory name, function is TO DO\n",
    "    def __init__(\n",
    "        self,\n",
    "        files, \n",
    "        X=X):\n",
    "        # print(\"this is the new and imporcces\")\n",
    "        \n",
    "        self.video_extractor = VideoDataMIP(files)\n",
    "\n",
    "        # list of dictionaries, each dictionary is one cell tracked throughout a video\n",
    "        #each cell has 3 data types in a dictionary\n",
    "            #mask, box and patches\n",
    "\n",
    "        for i in files:\n",
    "            self.video_extractor.extract_all_traces(i[1], X+1)\n",
    "        \n",
    "        \n",
    "        self.cell_dict = []\n",
    "\n",
    "        for key in self.video_extractor.data:\n",
    "            entry = self.video_extractor.data[key][\"traces\"]\n",
    "            for cell in entry:\n",
    "                patches = [np.array(p) for p in cell[\"patches\"]]\n",
    "                boxes = [np.array(b) for b in cell['boxes']]\n",
    "                masks = [np.array(m) for m in cell['masks']]\n",
    "                \n",
    "                self.cell_dict.append((boxes, masks, patches)) #cell dict is a list of 3 types by sequence\n",
    "\n",
    "        self.num_cells = len(self.cell_dict) #this is a list of how many sequences we have\n",
    "              \n",
    "    def __len__(self):\n",
    "        return self.num_cells\n",
    "        \n",
    "\n",
    "    def get_centroids(self, boxes, masks):\n",
    "        N = len(masks)\n",
    "        res = []\n",
    "        centroids = [skm.centroid(binary.astype(np.uint8)) for binary in masks]\n",
    "        # print(\"centroids are\", centroids)\n",
    "        for i in range(N):\n",
    "            c = centroids[i]\n",
    "            ymin, xmin = boxes[i][:2]\n",
    "            # print(\"didnt even get to res\")\n",
    "            res.append([xmin+c[0], ymin+c[1]])\n",
    "            # print(\"results is \", type(res), \"shape\", np.array(res).shape)\n",
    "        return(np.array(res) - res[0]) \n",
    "   \n",
    "    def pad_arrays(self, array, pad_amt=200):\n",
    "        # print(\"array shape 0 is\", array.shape[0])\n",
    "        # print(\"array shape 1\", array.shape[1])\n",
    "\n",
    "\n",
    "        # print(\"x would be\", pad_amt-array.shape[0])\n",
    "        # print(\"y would be\", pad_amt-array.shape[1])\n",
    "        pad_width = ((0, pad_amt - array.shape[0]), (0, pad_amt - array.shape[1]))\n",
    "        # print(\"pad width is \",pad_width)\n",
    "\n",
    "        padded_array = np.pad(array, pad_width, mode='constant')\n",
    "        # print(\"pad array is is \",padded_array.shape)\n",
    "        return padded_array\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        cell_sequences = self.cell_dict[idx]  #this is the first sequence of 10 cells\n",
    "        boxes = cell_sequences[0]\n",
    "        masks = cell_sequences[1]\n",
    "        patches = cell_sequences[2]\n",
    "\n",
    "\n",
    "        for cell_mask_num in np.arange(len(masks)): #should be sequence length (10) masks\n",
    "                \n",
    "                #make masks 1 and 0s\n",
    "                cell_time = np.array(masks[cell_mask_num], dtype=np.int32)\n",
    "                cell_time = np.where(cell_time >= 0, cell_time, 1)\n",
    "\n",
    "                #pad masks\n",
    "                # print(\"cell_time shape, before padding\", cell_time.shape)\n",
    "                cell_time = self.pad_arrays(cell_time)\n",
    "                # print(\"cell_time shape, after padding\", cell_time.shape)\n",
    "                masks[cell_mask_num] = cell_time\n",
    "\n",
    "                #pad patches\n",
    "                # print(\"cell_time PATCH shape, before padding\", cell_time.shape)\n",
    "                cell_time_patch = np.array(patches[cell_mask_num], dtype=np.int32)\n",
    "                # print(\"cell_time PATCH shape, after padding\", cell_time.shape)\n",
    "\n",
    "                cell_time_patch = self.pad_arrays(cell_time_patch)\n",
    "\n",
    "                patches[cell_mask_num] = cell_time_patch\n",
    "\n",
    "        # print(boxes)\n",
    "\n",
    "        centroids = self.get_centroids(boxes, masks)\n",
    "    def __getitem__(self, idx):\n",
    "        #should T dictiaries for . \n",
    "        #If any cells don't exist (ie we look for index n, n+1, n+2, n+3..) and the list is of \n",
    "        #length n+2, everything that doesn't exist (n+2, n+3...) should just be self.end, our \"end token\"\n",
    "    \n",
    "        #turn start and cell\n",
    "        # print(self.list_timesteps_start, self.list_timesteps_end)\n",
    "\n",
    "        # cell_idx, start = self.binary_search(self.list_timesteps_end, self.list_timesteps_start, idx)\n",
    "\n",
    "        cell = self.cell_dict[idx]\n",
    "\n",
    "        boxes = cell[\"boxes\"]\n",
    "        # masks = cell[\"masks\"]\n",
    "        patches = cell[\"patches\"]\n",
    "\n",
    "\n",
    "        #make the masks all 1's for all the cells\n",
    "        # for cell_mask in np.arange(len(masks)): #there are 67 cells, each cell is a list of 10 cells\n",
    "        #     for cell_time in masks[cell_mask]:\n",
    "        #         cell_time = np.array(cell_time, dtype=np.int32)\n",
    "        #         cell_time = np.where(cell_time < 0, cell_time, 1)\n",
    "        #         masks[cell_mask] = cell_time\n",
    "\n",
    "        # box = boxes[start]\n",
    "        # mask = masks[start]\n",
    "        # patch = patches[start]\n",
    "\n",
    "        # # Get the dictionaries for the next T cells\n",
    "        # next_T_frames_box = []\n",
    "        # next_T_frames_masks = []\n",
    "        # next_T_frames_patches = []\n",
    "        # for i in range(1, self.T + 1):\n",
    "        #     # print(\"Getting next frame\", i)\n",
    "        #     if start + i < self.list_timesteps_end[cell_idx]-self.list_timesteps_start[cell_idx]:\n",
    "        #         next_T_frames_box.append(boxes[start + i])\n",
    "        #         next_T_frames_masks.append(masks[start+1].flatten())\n",
    "        #         #change all the numbers to smth else --> do do mask\n",
    "        #         next_T_frames_patches.append(patches[start+i].flatten())\n",
    "            \n",
    "        \n",
    "        # print(\"shoudl be 10\", len(patches))\n",
    "        # print(\"after all, should be 250,250\", patches[0].shape)\n",
    "\n",
    "        # print(\"shoudl be 10 MASKSS\", len(masks))\n",
    "        # print(\"after all, should be 250,250 MASKS\", masks[0].shape)\n",
    "\n",
    "        return centroids, masks, patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torch.utils.data import random_split\n",
    "\n",
    "# root_dir = \"/mnt/datadisk/\"\n",
    "# subdirs = ['FactinProcessed', 'FactinMIP']\n",
    "\n",
    "# mip_video_files = [\n",
    "#     ('mip', 3),\n",
    "#     ('mip', 6),\n",
    "#     ('mip', 9),\n",
    "# ]\n",
    "\n",
    "# mip_test_files = [('mip', 3)] #so that I dont keep crashing the kernel\n",
    "\n",
    "# # file = root_dir + subdirs\n",
    "\n",
    "# video_extractor = VideoDataMIP(mip_test_files)\n",
    "\n",
    "#         # list of dictionaries, each dictionary is one cell tracked throughout a video\n",
    "#         #each cell has 3 data types in a dictionary\n",
    "#             #mask, box and patches\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# files = [('mip', 3)]\n",
    "\n",
    "# video_extractor = VideoDataMIP(files)\n",
    "\n",
    "# # list of dictionaries, each dictionary is one cell tracked throughout a video\n",
    "# #each cell has 3 data types in a dictionary\n",
    "# #mask, box and patches\n",
    "# X=10\n",
    "\n",
    "# for i in files:\n",
    "#     video_extractor.extract_all_traces(i[1], X)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import skimage.measure as skm\n",
    "        \n",
    "        \n",
    "# cell_dict = []\n",
    "\n",
    "# for key in video_extractor.data:\n",
    "#             entry = video_extractor.data[key][\"traces\"]\n",
    "#             for cell in entry:\n",
    "#                 # print(\"cell is of type\", type(cell))\n",
    "#                 sequence = []\n",
    "#                 patches = [np.array(p) for p in cell[\"patches\"]]\n",
    "#                 boxes = [np.array(b) for b in cell['boxes']]\n",
    "#                 masks = [np.array(m) for m in cell['masks']]\n",
    "#                 # print(\"first patch\", patches[0])\n",
    "#                 # print(\"first patch as arraty\", np.array(patches[0]))\n",
    "#                 # print(\"first patch shape\", np.array(patches[0]).shape)\n",
    "\n",
    "#                 # print(\"first mask\", masks[0])\n",
    "#                 # print(\"first mask as arraty\", np.array(masks[0]))\n",
    "#                 # print(\"first mask SHAPPEPEEEE\", np.array(masks[0]).shape)\n",
    "\n",
    "#                 # for i in np.arange(len(patches)):\n",
    "#                 # #     sequence.append({\n",
    "#                 # #         \"boxes\": boxes[i],\n",
    "#                 # #         \"patches\": patches[i],\n",
    "#                 # #         \"masks\": masks[i]\n",
    "#                 # #     })\n",
    "#                 # #     print(sequence)\n",
    "#                 #         sequence.append((boxes, patches, masks))\n",
    "#                 #         # print(\"sequence should be ten lists of 3\",len(sequence))\n",
    "#                 #         # print(\"here is the first list of ten, should be 3 long\", len(sequence[0]))\n",
    "#                 cell_dict.append((boxes, masks, patches)) #cell dict is a list of 3 types by sequence\n",
    "\n",
    "# num_cells = len(cell_dict) #this is a list of how many sequences we have\n",
    "\n",
    "\n",
    "# cell_sequences = cell_dict[0]  #this is the first sequence of 10 cells\n",
    "# boxes = cell_sequences[0]\n",
    "# patches = cell_sequences[1]\n",
    "# masks = cell_sequences[2]\n",
    "\n",
    "\n",
    "# def get_centroids(self, cell_data):\n",
    "#         boxes, masks = np.array(cell_data['boxes']), cell_data['masks']\n",
    "#         N = len(masks)\n",
    "#         res = []\n",
    "#         centroids = [skm.centroid(binary.astype(np.uint8)) for binary in masks]\n",
    "#         for i in range(N):\n",
    "#             c = centroids[i]\n",
    "#             ymin, xmin = boxes[i][:2]\n",
    "#             res.append([xmin+c[0], ymin+c[1]])\n",
    "#         return(np.array(res) - res[0,:]) \n",
    "   \n",
    "# def pad_arrays(self, array, pad_amt=200):\n",
    "#         print(pad_amt)\n",
    "#         print(\"tyep\", type(array))\n",
    "#         print(\"array shape 0 is\", array.shape[0])\n",
    "#         print(\"array shape 1\", array.shape[1])\n",
    "\n",
    "#         print(\"x would be\", pad_amt-array.shape[0])\n",
    "#         print(\"y would be\", pad_amt-array.shape[1])\n",
    "#         pad_width = ((0, pad_amt - array.shape[0]), (0, pad_amt - array.shape[1]))\n",
    "#         print(pad_width)\n",
    "\n",
    "#         padded_array = np.pad(array, pad_width, mode='constant')\n",
    "#         return padded_array\n",
    "\n",
    "\n",
    "# #make the masks all 1's for all the cells\n",
    "# for cell_mask_num in np.arange(len(masks)): #should be sequence length (10) masks\n",
    "#         # print(\"cell_mask_num\", cell_mask_num)\n",
    "#         cell_time = np.array(masks[cell_mask_num], dtype=np.int32)\n",
    "#         # print(\"this should be one mask\", cell_time)\n",
    "#         cell_time = np.where(cell_time >= 0, cell_time, 1)\n",
    "#         # print(\"min and max\", np.min(cell_time), np.max(cell_time))\n",
    "#         print(\"cell time shape and cell time 1 \", cell_time.shape)\n",
    "#         cell_time = pad_arrays(cell_time)\n",
    "#         print(\"cell time shape and cell time 2\", cell_time.shape)\n",
    "\n",
    "#         masks[cell_mask_num] = cell_time\n",
    "\n",
    "\n",
    "#         cell_time_patch = np.array(masks[cell_mask_num], dtype=np.int32)\n",
    "#         cell_time_patch = pad_arrays(cell_time_patch)\n",
    "\n",
    "#         masks[cell_mask_num] = cell_time_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# masks = [v['masks'] for v in video_extractor.data[3]['traces']]#this is a list of 10 masks\n",
    "\n",
    "\n",
    "# import numpy as np\n",
    "\n",
    "# def pad_arrays(array, padding_amt=200):\n",
    "#     pad_width = ((0, 200 - array.shape[0]), (0, 200 - array.shape[1]))\n",
    "#     padded_array = np.pad(array, pad_width, mode='constant')\n",
    "#     return padded_array\n",
    "\n",
    "\n",
    "# #get make all the values == 1\n",
    "\n",
    "# #padding\n",
    "# shapes = []\n",
    "# for cell_mask in masks: #there are 67 cells, each cell is a list of 10 cells\n",
    "#     for cell_time in cell_mask:\n",
    "#         cell_time = np.array(cell_time, dtype=np.int32)\n",
    "#         cell_time = np.where(cell_time < 0, cell_time, 1)\n",
    "#         # print(cell_time.shape)\n",
    "        \n",
    "#         cell_time = pad_arrays(cell_time)\n",
    "#         shapes.append(cell_time.shape)\n",
    "#         #gonna pad to [180, 180], but can be changed after\n",
    "\n",
    "    \n",
    "# print(np.max(shapes))\n",
    "# print(np.min(shapes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cell_time = np.array([[1, 2, 3], [4, 5, 6]])\n",
    "\n",
    "# print(np.max(cell_time))\n",
    "\n",
    "# cell_time = np.where(cell_time < 1, cell_time, 1)\n",
    "\n",
    "\n",
    "# print(np.max(cell_time))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from skimage.measure import centroid\n",
    "# files = [('mip', 3)]\n",
    "# video_extractor = VideoDataMIP(files)\n",
    "\n",
    "# # test_patches = video_extractor.data[3]['traces'][0]['patches'] #this is a list of X paddings\n",
    "\n",
    "# # #get centroids\n",
    "\n",
    "# # centroids = ([centroid(t) for t in test_patches])\n",
    "# # centroids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading in MIP 3\n",
      "Loading dicty_factin_pip3-03_MIP.czi with dims [{'X': (0, 474), 'Y': (0, 2048), 'C': (0, 2), 'T': (0, 90)}]\n",
      "frames 3: (90, 2048, 474)\n",
      "Extracting traces from 0:11\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n",
      "Extracting traces from 11:22\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n",
      "Extracting traces from 22:33\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n",
      "Extracting cell  7\n",
      "Extracting cell  8\n",
      "Extracting traces from 33:44\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n",
      "Extracting cell  7\n",
      "Extracting traces from 44:55\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n",
      "Extracting traces from 55:66\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting traces from 66:77\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting traces from 77:88\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  0\n",
      "Extracting cell  1\n",
      "Extracting cell  2\n",
      "Extracting cell  3\n",
      "Extracting cell  4\n",
      "Extracting cell  5\n",
      "Extracting cell  6\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import random_split\n",
    "\n",
    "mip_video_files = [\n",
    "    ('mip', 3)\n",
    "]\n",
    "\n",
    "dataset = CellBoxMaskPatch(mip_video_files, 10) # file, S, T\n",
    "\n",
    "train, eval, test = random_split(dataset, [0.4, 0.2, 0.4])\n",
    "\n",
    "input_datasets = {}\n",
    "input_datasets[\"train\"] = train\n",
    "input_datasets[\"eval\"] = eval\n",
    "input_datasets[\"test\"] = test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # timestep = [[0] for t in range(T)]\n",
    "# # timestep\n",
    "\n",
    "# # batch = ((np.array([1 , 2, 3]), 1, 2) , 3), ((np.array([1 , 2, 3]), 1, 2) , 3)\n",
    "# # current_boxes = torch.stack([torch.tensor(b[0][0], dtype=torch.long) for b in batch], dim=0)\n",
    "# # current_boxes\n",
    "#         #current     #next   #boxes                     #masks\n",
    "# batch  = [[1, 2, 3], [[[1,2, 3, 4], [1, 2, 3,4]], [[1, 2, 3], [1, 2, 3]], [[1, 2, 3], [1,2, 3]]]], [[1, 2, 3], [[[1,2, 3, 4], [1, 2, 3,4]], [[1, 2, 3], [1, 2, 3]], [[1, 2, 3], [1,2, 3]]]]\n",
    "\n",
    "\n",
    "\n",
    "# next_boxes = torch.stack([torch.tensor(b[1][0], dtype=torch.long) for b in batch], dim=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(batch, mode_box, mode_mask, mode_patch):\n",
    "    \n",
    "\n",
    "    centroids = torch.stack([torch.tensor(b[0], dtype=torch.int) for b in batch], dim=0)\n",
    "    print(centroids.shape)\n",
    "    current_masks = torch.stack([torch.tensor(b[1], dtype=torch.long) for b in batch], dim=0)\n",
    "    current_masks = current_masks.reshape([4, 10, 40000])\n",
    "    print(current_masks.shape)\n",
    "    current_patches = torch.stack([torch.tensor(b[2], dtype=torch.long) for b in batch], dim=0)\n",
    "    current_patches = current_patches.reshape([4, 10, 40000])\n",
    "\n",
    "    print(current_patches.shape)\n",
    "\n",
    "    selected_tensors = []\n",
    "    if mode_box:\n",
    "        selected_tensors.append(centroids)\n",
    "    if mode_mask:\n",
    "        selected_tensors.append(current_masks)\n",
    "    if mode_patch:\n",
    "        selected_tensors.append(current_patches)\n",
    "\n",
    "    # Concatenate selected tensors along the last dimension\n",
    "    combined_tensor = torch.cat(selected_tensors, dim=-1)\n",
    "\n",
    "    # Reshape to add the singleton dimension\n",
    "    combined_tensor = combined_tensor #.unsqueeze(1)\n",
    "\n",
    "    return combined_tensor, centroids\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode_box = True\n",
    "mode_mask = True\n",
    "mode_patch = False\n",
    "\n",
    "dataloaders = {}\n",
    "dataloaders['train'] = torch.utils.data.DataLoader(\n",
    "    input_datasets['train'],\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    "    collate_fn=lambda batch: collate_fn(batch, mode_box, mode_mask, mode_patch)\n",
    ")\n",
    "\n",
    "dataloaders['test'] = torch.utils.data.DataLoader(\n",
    "    input_datasets['test'],\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    "    collate_fn=lambda batch: collate_fn(batch, mode_box, mode_mask, mode_patch)\n",
    ")\n",
    "\n",
    "dataloaders['eval'] = torch.utils.data.DataLoader(\n",
    "    input_datasets['eval'],\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    num_workers=0,\n",
    "    collate_fn=lambda batch: collate_fn(batch, mode_box, mode_mask, mode_patch)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 10, 2])\n",
      "torch.Size([4, 10, 40000])\n",
      "torch.Size([4, 10, 40000])\n",
      "Input: torch.Size([4, 10, 40002]) Centroids torch.Size([4, 10, 2])\n"
     ]
    }
   ],
   "source": [
    "for batch in dataloaders['train']:\n",
    "    # input, output = batch\n",
    "    # Process the batches as needed\n",
    "    # print(\"Input Shape:\", input.shape)\n",
    "    # print(\"Output Shape:\", output.shape)\n",
    "    print(\"Input:\", batch[0].shape, \"Centroids\", batch[1].shape)\n",
    "    \n",
    "    break  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create an LSTM RNN\n",
    "\n",
    "#to create a pytorch LSTM \n",
    "#The first axis is the sequence itself\n",
    "# the second indexes instances in the mini-batch\n",
    "# third indexes elements of the input. \n",
    "\n",
    "batch_size = 4\n",
    "\n",
    "class LSTM(nn.Module):\n",
    "    def __init__(self, input_size, hidden_dim, num_layers=2):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.lstm = nn.LSTM(input_size, hidden_dim, num_layers=2, batch_first=True) #stacking 2 LSTMs\n",
    "\n",
    "        # The linear layer that maps from hidden state space to output space\n",
    "        # Currently set equal to 2 for the x & y dimension of the centroid\n",
    "        self.fc = nn.Linear(hidden_dim, 2)\n",
    "    def forward(self, input):\n",
    "\n",
    "\n",
    "        #h_0: tensor of shape (D∗num_layers,N,Hout​) containing the initial hidden\n",
    "        # state for each element in the input sequence. Defaults to zeros if (h_0, c_0) is not provided.\n",
    "        h0 = torch.zeros(self.num_layers, input.size(0), self.hidden_dim).to(input.device)\n",
    "        c0 = torch.zeros(self.num_layers, input.size(0), self.hidden_dim).to(input.device)\n",
    "\n",
    "        # We need to detach as we are doing truncated backpropagation through time (BPTT)\n",
    "        # If we don't, we'll backprop all the way to the start even after going through another batch\n",
    "        out, _ = self.lstm(input, (h0.detach(), c0.detach()))\n",
    "\n",
    "        # Predict coordinate + bound to 0-200.\n",
    "        out = self.fc(out)\n",
    "        final = out[:,-1,:]\n",
    "        out = torch.sigmoid(out) * 200\n",
    "        return out\n",
    "        # lstm_out, h_n = self.lstm(input, batch_size, -1) #output and final hidden state\n",
    "        # output = self.hidden2out(lstm_out.view(len(input), -1))\n",
    "        # #output = F.log_softmax(output, dim=1)\n",
    "        # return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.2055,  0.6472],\n",
      "         [ 0.4410, -0.7373],\n",
      "         [-0.5874,  0.1577],\n",
      "         [ 0.2047,  0.1498],\n",
      "         [-0.3650,  0.2788],\n",
      "         [-0.3681,  0.7801],\n",
      "         [-0.2030,  0.6843],\n",
      "         [-1.4174, -0.9384],\n",
      "         [ 1.1115, -0.1738],\n",
      "         [ 1.2157,  0.7619]],\n",
      "\n",
      "        [[ 0.2684,  0.0097],\n",
      "         [-1.4432, -0.4194],\n",
      "         [-0.9439, -0.3813],\n",
      "         [-0.4078, -0.9100],\n",
      "         [-1.4686,  0.9237],\n",
      "         [ 3.4683, -0.1606],\n",
      "         [-0.4892, -0.9693],\n",
      "         [ 0.0670,  2.4959],\n",
      "         [ 0.5985,  0.5135],\n",
      "         [ 0.5734,  0.2431]]])\n",
      "tensor([[1.2157, 0.7619],\n",
      "        [0.5734, 0.2431]])\n",
      "torch.Size([2, 10, 2])\n"
     ]
    }
   ],
   "source": [
    "# dummy_input_data = torch.randn(2, 10, input_size)\n",
    "# print(dummy_input_data)\n",
    "# print(dummy_input_data[:,-1,:])\n",
    "# print(dummy_input_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 10, 2])\n",
      "tensor([[[97.7781, 98.6298],\n",
      "         [97.7656, 98.6439],\n",
      "         [97.7639, 98.8019],\n",
      "         [97.7323, 98.9574],\n",
      "         [97.7105, 99.2236],\n",
      "         [97.6879, 99.2642],\n",
      "         [97.7201, 99.2554],\n",
      "         [97.6608, 99.1939],\n",
      "         [97.6966, 99.1502],\n",
      "         [97.6950, 98.9583]],\n",
      "\n",
      "        [[97.8083, 98.3526],\n",
      "         [97.7845, 98.2537],\n",
      "         [97.6489, 98.3474],\n",
      "         [97.6066, 98.4637],\n",
      "         [97.6593, 98.3532],\n",
      "         [97.6939, 98.4748],\n",
      "         [97.6884, 98.6090],\n",
      "         [97.6486, 98.6024],\n",
      "         [97.6121, 98.8213],\n",
      "         [97.6308, 98.7603]],\n",
      "\n",
      "        [[97.7912, 98.6616],\n",
      "         [97.7845, 98.6935],\n",
      "         [97.7111, 98.6017],\n",
      "         [97.7026, 98.3273],\n",
      "         [97.6435, 98.2857],\n",
      "         [97.6460, 98.3809],\n",
      "         [97.7210, 98.4604],\n",
      "         [97.7512, 98.5766],\n",
      "         [97.7530, 98.6088],\n",
      "         [97.7554, 98.6745]],\n",
      "\n",
      "        [[97.7983, 98.4820],\n",
      "         [97.7394, 98.7360],\n",
      "         [97.7440, 98.8022],\n",
      "         [97.7574, 98.7472],\n",
      "         [97.7416, 98.8872],\n",
      "         [97.6682, 98.9912],\n",
      "         [97.6139, 98.9712],\n",
      "         [97.6035, 98.8472],\n",
      "         [97.6157, 98.5863],\n",
      "         [97.5958, 98.6039]]], grad_fn=<MulBackward0>)\n",
      "torch.Size([4, 10, 2])\n"
     ]
    }
   ],
   "source": [
    "input_size = 2 # this is centroid input only\n",
    "hidden_size = 100\n",
    "num_layers = 2\n",
    "epochs = 1\n",
    "sequence_length = 10 #how many frames we process per input\n",
    "\n",
    "model = LSTM(input_size, hidden_size, num_layers)\n",
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "dummy_input_data = torch.randn(batch_size, 10, input_size)\n",
    "\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    total_correct = 0\n",
    "    for batch in dataloaders['train']:\n",
    "        optimizer.zero_grad()\n",
    "        inputs, outputs = batch[0], batch[1]\n",
    "        inputs, outputs = inputs.to(device), outputs.to(device)\n",
    "        pred = model(inputs[:, :sequence_length-1, :])\n",
    "        total_correct += torch.sum(pred == outputs.data)\n",
    "        loss = criterion(pred, outputs[:,-1,:])\n",
    "        total_loss += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f\"training loss: {total_loss / len(dataloaders['train'])}, training accuracy: {total_correct / len(dataloaders['train'])}\")\n",
    "    return model\n",
    "\n",
    "def eval():\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    total_correct = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in dataloaders['dev']:\n",
    "            inputs, outputs = batch[0], batch[1]\n",
    "            inputs, outputs = inputs.to(device), outputs.to(device)\n",
    "            pred = model(inputs[:, :sequence_length-1, :])\n",
    "            total_correct += torch.sum(pred == outputs.data)\n",
    "            loss = criterion(pred, outputs[:,-1,:])\n",
    "            total_loss += loss.item()\n",
    "    print(f\"validation loss: {total_loss / len(dataloaders['train'])}, validation accuracy: {total_correct / len(dataloaders['train'])}\")\n",
    "    return total_correct / len(dataloaders['train']\n",
    "\n",
    "def train_model():\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "    for epoch in num_epochs:\n",
    "        train()\n",
    "        curr_acc = eval()\n",
    "        if curr_acc > best_acc:\n",
    "            best_acc = epoch_acc\n",
    "            best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'input_dim' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[20], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m T \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m10\u001b[39m \u001b[38;5;66;03m#len predictions\u001b[39;00m\n\u001b[1;32m      9\u001b[0m hidden_dims \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m100\u001b[39m\n\u001b[0;32m---> 10\u001b[0m output_dim \u001b[38;5;241m=\u001b[39m \u001b[43minput_dim\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msigmoid\u001b[39m(x):\n\u001b[1;32m     13\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;241m1\u001b[39m\u001b[38;5;241m/\u001b[39m \u001b[38;5;241m1\u001b[39m\u001b[38;5;241m+\u001b[39m np\u001b[38;5;241m.\u001b[39mexp(\u001b[38;5;241m-\u001b[39mx)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'input_dim' is not defined"
     ]
    }
   ],
   "source": [
    "#mip only had one zplane \n",
    "# Example input data\n",
    "#parameters for mask and patch\n",
    "data_dim =  1000 #TO DO: find the biggest box and set it to this\n",
    "learning_rate = 0.0001\n",
    "epochs = 25\n",
    "\n",
    "T = 10 #len predictions\n",
    "hidden_dims = 100\n",
    "output_dim = input_dim\n",
    "\n",
    "def sigmoid(x):\n",
    "    return 1/ 1+ np.exp(-x)\n",
    "\n",
    "batch_size = 3\n",
    "seq_length = 4\n",
    "input_dim = 4\n",
    "hidden_dim = 5\n",
    "num_layers = 1\n",
    "#input is (batch, seq, feature)\n",
    "input_data = torch.randn(batch_size, input_dim, input_dim)\n",
    "\n",
    "# Create LSTM model\n",
    "lstm_model = LSTM(input_dim, hidden_dim, output_dim, num_layers)\n",
    "\n",
    "# Forward pass with desired sequence length T\n",
    "T = seq_length\n",
    "output = lstm_model(input_data)\n",
    "\n",
    "print(\"Input shape:\", input_data.shape)\n",
    "print(\"Output shape:\", output.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "input_dim = 0\n",
    "if mode_box == True:\n",
    "    input_dim +=  4 #(T, 4) for the bounding boxes\n",
    "if mode_mask == True:\n",
    "    input_dim += box_shape.flatten #flattened version of [H, W]\n",
    "if mode_patch == True:\n",
    "    input_dim += box_shape.flatten #flattened version of [H, W]\n",
    "\n",
    "learning_rate = 0.0001\n",
    "epochs = 25\n",
    "T = 2\n",
    "hidden_dims = 100\n",
    "output_dim = input_dim\n",
    "batch_size = 1\n",
    "model = LSTM(input_dim, hidden_dims, output_dim)\n",
    "loss_function = nn.NLLLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do your training loop + eval, etc...\n",
    "def train_model(model, criterion, optimizer, scheduler, num_epochs=epochs):\n",
    "    since = time.time()\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print(f'Epoch {epoch}/{num_epochs - 1}')\n",
    "        print('-' * 10)\n",
    "\n",
    "        # Each epoch has a training and validation phase\n",
    "        for phase in ['train', 'dev']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # Set model to training mode\n",
    "            else:\n",
    "                model.eval()   # Set model to evaluate mode\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "\n",
    "            # Iterate over data.\n",
    "            for inputs, labels in dataloaders[phase]:  #inputs shoudl be [B, S, D] #labes be a [B, T, D] batch, time, sequence shape\n",
    "                inputs = inputs.to(device)\n",
    "\n",
    "                 # zero the parameter gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # forward\n",
    "                # track history if only in train\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs, labels) #B, T, D\n",
    "\n",
    "                    #LOSS AND PRED EDIT STUFF HERE AS WE TALKED ABOUT\n",
    "\n",
    "                    # backward + optimize only if in training phase\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += torch.sum(preds == labels.data)\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / dataset_sizes[phase]\n",
    "            epoch_acc = running_corrects.double() / dataset_sizes[phase]\n",
    "\n",
    "            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')\n",
    "\n",
    "            # deep copy the model\n",
    "            if phase == 'dev' and epoch_acc > best_acc:\n",
    "                best_acc = epoch_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')\n",
    "    print(f'Best val Acc: {best_acc:4f}')\n",
    "\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    return model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lukeenv",
   "language": "python",
   "name": "lukeenv"
  },
  "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}