debug_relationformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"2\"\n",
    "import torch\n",
    "import glob\n",
    "import math\n",
    "import numpy as np\n",
    "from medpy.io import load\n",
    "import pyvista\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from dataset_vessel3d import build_vessel_data\n",
    "from monai.data import DataLoader\n",
    "from torch import nn\n",
    "from skimage.morphology import skeletonize_3d\n",
    "from skimage.measure import marching_cubes_lewiner\n",
    "from scipy.sparse import csr_matrix\n",
    "from models.detr_transformer_3D import build_detr_transformer\n",
    "from losses import SetCriterion\n",
    "from models import build_model\n",
    "import networkx as nx\n",
    "# %matplotlib widget\n",
    "import yaml\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "import json\n",
    "from models.position_encoding import PositionalEncoding3D\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from scipy import ndimage\n",
    "import open3d as o3d\n",
    "from inference import relation_infer\n",
    "from scipy.sparse.csgraph import connected_components\n",
    "from utils import save_input, save_output\n",
    "from utils import Bresenham3D\n",
    "import pyvista as pv\n",
    "import pdb\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_test_sample(image, points, edges):\n",
    "    meshes = []\n",
    "    graphs = []\n",
    "    image = image[:,54:-54,54:-54]\n",
    "    porder_points = [\n",
    "        [0, 0, 0],\n",
    "        [1, 0, 0],\n",
    "        [0, 1, 0],\n",
    "        [1, 1, 0],\n",
    "        [0, 0, 1],\n",
    "        [1, 0, 1],\n",
    "        [0, 1, 1],\n",
    "        [1, 1, 1],\n",
    "    ]\n",
    "    border_edges = [\n",
    "        [0, 1],\n",
    "        [0, 2],\n",
    "        [1, 3],\n",
    "        [2, 3],\n",
    "        [4, 5],\n",
    "        [4, 6],\n",
    "        [5, 7],\n",
    "        [6, 7],\n",
    "        [0, 4],\n",
    "        [1, 5],\n",
    "        [2, 6],\n",
    "        [3, 7],\n",
    "    ]\n",
    "\n",
    "    # edges = np.concatenate((np.int32(2*np.ones((edges.shape[0],1))), edges), 1)\n",
    "    # gt_graph = pyvista.PolyData(points)\n",
    "    # gt_graph.lines = edges.flatten()\n",
    "    ref_color = [[1, 0, 0] for i in range(len(edges))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(points + np.array([0.6, 0, 0])),\n",
    "                                lines=o3d.utility.Vector2iVector(edges),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    graphs.append(ref_line_set)\n",
    "    point_cloud = o3d.geometry.PointCloud()\n",
    "    point_cloud.points = o3d.utility.Vector3dVector(points + np.array([0.6, 0, 0]))\n",
    "    point_cloud.paint_uniform_color([0, 1, 0])\n",
    "    graphs.append(point_cloud)\n",
    "    \n",
    "    ref_color = [[0.2, 0.2, 0.2] for i in range(len(border_edges))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(porder_points + np.array([0.6, 0, 0])),\n",
    "                                lines=o3d.utility.Vector2iVector(border_edges),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    graphs.append(ref_line_set)\n",
    "    \n",
    "    verts, faces, norms, vals = marching_cubes_lewiner(image>0.0, level=0)\n",
    "    verts = verts/np.array(image.shape) \n",
    "\n",
    "    mesh = np.concatenate((faces[:,:2], faces[:,1:]), axis=0)\n",
    "    adjucency = np.zeros((verts.shape[0], verts.shape[0]))\n",
    "    \n",
    "    for e in mesh:\n",
    "        adjucency[e[0], e[1]] =1.0\n",
    "        adjucency[e[1], e[0]] =1.0\n",
    "    \n",
    "    adjucency = np.triu(adjucency)\n",
    "    mesh = np.array(np.where(np.triu(adjucency)>0)).T\n",
    "    # mesh = np.concatenate((np.int32(2*np.ones((mesh.shape[0],1))), mesh), 1)\n",
    "    # gt_mesh = pyvista.PolyData(verts)\n",
    "    # gt_mesh.lines = mesh.flatten()\n",
    "    pred_color = [[0, 0, 1] for i in range(len(mesh))]\n",
    "    pred_line_set = o3d.geometry.LineSet(\n",
    "                               points=o3d.utility.Vector3dVector(verts- np.array([0.6, 0, 0])),\n",
    "                               lines=o3d.utility.Vector2iVector(mesh),\n",
    "                           )\n",
    "    \n",
    "    pred_line_set.colors = o3d.utility.Vector3dVector(pred_color)\n",
    "    meshes.append(pred_line_set)\n",
    "    \n",
    "    pred_color = [[0.2, 0.2, 0.2] for i in range(len(border_edges))]\n",
    "    pred_line_set = o3d.geometry.LineSet(\n",
    "                               points=o3d.utility.Vector3dVector((porder_points - np.array([0.6, 0, 0]))),\n",
    "                               lines=o3d.utility.Vector2iVector(border_edges),\n",
    "                           )\n",
    "    \n",
    "    pred_line_set.colors = o3d.utility.Vector3dVector(pred_color)\n",
    "    meshes.append(pred_line_set)\n",
    "    \n",
    "    \n",
    "    o3d.visualization.draw_geometries(meshes+graphs)\n",
    "    \n",
    "    # plotter = pv.Plotter(shape=(1, 2))\n",
    "\n",
    "    # plotter.subplot(0, 0)\n",
    "    # plotter.add_mesh(gt_mesh, show_edges=True, color=\"b\")\n",
    "\n",
    "    # plotter.subplot(0, 1)\n",
    "    # plotter.add_mesh(gt_graph, show_edges=True, color=\"r\")\n",
    "\n",
    "    # Display the window\n",
    "    # plotter.show()\n",
    "    \n",
    "\n",
    "def plot_val_rel_sample(image, points1, edges1, points2, edges2, attn_map=None, relative_coords=True):\n",
    "    ref_line_sets = []\n",
    "    pred_line_sets = []\n",
    "    porder_points = [\n",
    "        [0, 0, 0],\n",
    "        [1, 0, 0],\n",
    "        [0, 1, 0],\n",
    "        [1, 1, 0],\n",
    "        [0, 0, 1],\n",
    "        [1, 0, 1],\n",
    "        [0, 1, 1],\n",
    "        [1, 1, 1],\n",
    "    ]\n",
    "    border_edges = [\n",
    "        [0, 1],\n",
    "        [0, 2],\n",
    "        [1, 3],\n",
    "        [2, 3],\n",
    "        [4, 5],\n",
    "        [4, 6],\n",
    "        [5, 7],\n",
    "        [6, 7],\n",
    "        [0, 4],\n",
    "        [1, 5],\n",
    "        [2, 6],\n",
    "        [3, 7],\n",
    "    ]\n",
    "\n",
    "    # edges = np.concatenate((np.int32(2*np.ones((edges.shape[0],1))), edges), 1)\n",
    "    # gt_graph = pyvista.PolyData(points)\n",
    "    # gt_graph.lines = edges.flatten()\n",
    "    ref_color = [[1, 0, 0] for i in range(len(edges1))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(points1),\n",
    "                                lines=o3d.utility.Vector2iVector(edges1),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    ref_line_sets.append(ref_line_set)\n",
    "    point_cloud = o3d.geometry.PointCloud()\n",
    "    point_cloud.points = o3d.utility.Vector3dVector(points1)\n",
    "    point_cloud.paint_uniform_color([0, 1, 0])\n",
    "    ref_line_sets.append(point_cloud)\n",
    "\n",
    "    ref_color = [[0.2, 0.2, 0.2] for i in range(len(border_edges))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(porder_points),\n",
    "                                lines=o3d.utility.Vector2iVector(border_edges),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    ref_line_sets.append(ref_line_set)\n",
    "    \n",
    "    \n",
    "    ref_color = [[1, 0, 0] for i in range(len(edges2))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(points2 + np.array([1.1, 0, 0])),\n",
    "                                lines=o3d.utility.Vector2iVector(edges2),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    ref_line_sets.append(ref_line_set)\n",
    "    point_cloud = o3d.geometry.PointCloud()\n",
    "    point_cloud.points = o3d.utility.Vector3dVector(points2 + np.array([1.1, 0, 0]))\n",
    "    point_cloud.paint_uniform_color([0, 1, 0])\n",
    "    ref_line_sets.append(point_cloud)\n",
    "\n",
    "    ref_color = [[0.2, 0.2, 0.2] for i in range(len(border_edges))]\n",
    "    ref_line_set = o3d.geometry.LineSet(\n",
    "                                points=o3d.utility.Vector3dVector(porder_points + np.array([1.1, 0, 0])),\n",
    "                                lines=o3d.utility.Vector2iVector(border_edges),\n",
    "                            )\n",
    "    ref_line_set.colors = o3d.utility.Vector3dVector(ref_color)\n",
    "    ref_line_sets.append(ref_line_set)\n",
    "\n",
    "    verts, faces, norms, vals = marching_cubes_lewiner(image>0.0, level=0)\n",
    "    verts = verts/np.array(image.shape) \n",
    "\n",
    "    mesh = np.concatenate((faces[:,:2], faces[:,1:]), axis=0)\n",
    "    adjucency = np.zeros((verts.shape[0], verts.shape[0]))\n",
    "\n",
    "    for e in mesh:\n",
    "        adjucency[e[0], e[1]] =1.0\n",
    "        adjucency[e[1], e[0]] =1.0\n",
    "\n",
    "    adjucency = np.triu(adjucency)\n",
    "    mesh = np.array(np.where(np.triu(adjucency)>0)).T\n",
    "    # mesh = np.concatenate((np.int32(2*np.ones((mesh.shape[0],1))), mesh), 1)\n",
    "    # gt_mesh = pyvista.PolyData(verts)\n",
    "    # gt_mesh.lines = mesh.flatten()\n",
    "    pred_color = [[0, 0, 1] for i in range(len(mesh))]\n",
    "    pred_line_set = o3d.geometry.LineSet(\n",
    "                               points=o3d.utility.Vector3dVector(verts- np.array([1.1, 0, 0])),\n",
    "                               lines=o3d.utility.Vector2iVector(mesh),\n",
    "                           )\n",
    "\n",
    "    pred_line_set.colors = o3d.utility.Vector3dVector(pred_color)\n",
    "    pred_line_sets.append(pred_line_set)\n",
    "\n",
    "    pred_color = [[0.2, 0.2, 0.2] for i in range(len(border_edges))]\n",
    "    pred_line_set = o3d.geometry.LineSet(\n",
    "                               points=o3d.utility.Vector3dVector((porder_points - np.array([1.1, 0, 0]))),\n",
    "                               lines=o3d.utility.Vector2iVector(border_edges),\n",
    "                           )\n",
    "\n",
    "    pred_line_set.colors = o3d.utility.Vector3dVector(pred_color)\n",
    "    pred_line_sets.append(pred_line_set)\n",
    "\n",
    "    o3d.visualization.draw_geometries(ref_line_sets+pred_line_sets)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class obj:\n",
    "    def __init__(self, dict1):\n",
    "        self.__dict__.update(dict1)\n",
    "        \n",
    "def dict2obj(dict1):\n",
    "    return json.loads(json.dumps(dict1), object_hook=obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = \"./configs/synth_3D.yaml\"\n",
    "\n",
    "with open(config_file) as f:\n",
    "    config = yaml.load(f, Loader=yaml.FullLoader)\n",
    "\n",
    "config = dict2obj(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = './data/vessel_data/test_data/'\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = build_vessel_data(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_graph_collate(batch):\n",
    "    images = torch.cat([item_ for item in batch for item_ in item[0]], 0).contiguous()\n",
    "    segs = torch.cat([item_ for item in batch for item_ in item[1]], 0).contiguous()\n",
    "    points = [item_ for item in batch for item_ in item[2]]\n",
    "    edges = [item_ for item in batch for item_ in item[3]]\n",
    "    return [images, segs, points, edges]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mn_dataset_loader = DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=lambda x: tuple([y_.to(device) for y_ in x_] for x_ in image_graph_collate(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2 torch.Size([2, 1, 64, 64, 64]) torch.Size([27, 3]) 2\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Input array must be at least 2x2x2.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-18-51d5be575a6b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mseg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoints\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimages\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpoints\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0medges\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mplot_test_sample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mseg\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpoints\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medges\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-4-97bd5dd216c9>\u001b[0m in \u001b[0;36mplot_test_sample\u001b[0;34m(image, points, edges)\u001b[0m\n\u001b[1;32m     51\u001b[0m     \u001b[0mgraphs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mref_line_set\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m     \u001b[0mverts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfaces\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnorms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvals\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmarching_cubes_lewiner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m>\u001b[0m\u001b[0;36m0.0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlevel\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     54\u001b[0m     \u001b[0mverts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mverts\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/media/Drives/Data/venvs/deepro/lib/python3.6/site-packages/skimage/measure/_marching_cubes_lewiner.py\u001b[0m in \u001b[0;36mmarching_cubes_lewiner\u001b[0;34m(volume, level, spacing, gradient_direction, step_size, allow_degenerate, use_classic, mask)\u001b[0m\n\u001b[1;32m    264\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m     return _marching_cubes_lewiner(volume, level, spacing, gradient_direction,\n\u001b[0;32m--> 266\u001b[0;31m                                    step_size, allow_degenerate, use_classic, mask)\n\u001b[0m\u001b[1;32m    267\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    268\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/media/Drives/Data/venvs/deepro/lib/python3.6/site-packages/skimage/measure/_marching_cubes_lewiner.py\u001b[0m in \u001b[0;36m_marching_cubes_lewiner\u001b[0;34m(volume, level, spacing, gradient_direction, step_size, allow_degenerate, use_classic, mask)\u001b[0m\n\u001b[1;32m    278\u001b[0m         \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Input volume should be a 3D numpy array.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    279\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mvolume\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m2\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mvolume\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m2\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mvolume\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 280\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input array must be at least 2x2x2.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    281\u001b[0m     volume = np.ascontiguousarray(volume,\n\u001b[1;32m    282\u001b[0m                                   np.float32)  # no copy if not necessary\n",
      "\u001b[0;31mValueError\u001b[0m: Input array must be at least 2x2x2."
     ]
    }
   ],
   "source": [
    "for images, seg, points, edges in mn_dataset_loader:\n",
    "    images = torch.stack(images)\n",
    "    seg = torch.stack(seg)\n",
    "    print(len(points), images.shape, points[0].shape, len(edges))\n",
    "    plot_test_sample(seg[0].squeeze().cpu().numpy(), points[0].cpu().numpy(), edges[0].cpu().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debug Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(config)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_path = './trained_weights/3D_vessel_checkpoint_epoch=100.pt'\n",
    "checkpoint = torch.load(ckpt_path, map_location='cpu')\n",
    "missing_keys, unexpected_keys = model.load_state_dict(checkpoint['net'], strict=False)\n",
    "unexpected_keys = [k for k in unexpected_keys if not (k.endswith('total_params') or k.endswith('total_ops'))]\n",
    "if len(missing_keys) > 0:\n",
    "    print('Missing Keys: {}'.format(missing_keys))\n",
    "if len(unexpected_keys) > 0:\n",
    "    print('Unexpected Keys: {}'.format(unexpected_keys))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 1\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 2\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 3\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 4\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 5\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 6\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 7\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 8\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 9\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 10\n",
      "\u001b[1;33m[Open3D WARNING] GLFW Error: X11: The DISPLAY environment variable is missing\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] Failed to initialize GLFW\u001b[0;m\n",
      "\u001b[1;33m[Open3D WARNING] [DrawGeometries] Failed creating OpenGL window.\u001b[0;m\n",
      "Iteration: 11\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "iteration = 0\n",
    "for images, seg, points, edges in mn_dataset_loader:\n",
    "    images = torch.stack(images)-0.5\n",
    "    seg = torch.stack(seg)\n",
    "\n",
    "    h, out = model(images)\n",
    "    out = relation_infer(h.detach(), out, model, config.MODEL.DECODER.OBJ_TOKEN, config.MODEL.DECODER.RLN_TOKEN)\n",
    "\n",
    "    plot_val_rel_sample(seg[0].squeeze().cpu().numpy(), points[0].cpu().numpy(), edges[0].cpu().numpy(), out['pred_nodes'][0], out['pred_rels'][0])\n",
    "    \n",
    "    iteration = iteration+1\n",
    "    print('Iteration:',iteration)\n",
    "    if iteration>10:\n",
    "        break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.9 ('deepro')",
   "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.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "0a65b42259c001818d23a7dae5036e8364bfc846daa533e50e38f7a76759b9bb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}