notebooks/view_npz_files.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90ca8dc3",
   "metadata": {},
   "source": [
    "# View npz data\n",
    "This notebooks shows you how to visualize the data in npz files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8a588883",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/lambouj/anaconda3/envs/brepnet/bin/python\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "import sys\n",
    "print(sys.executable)\n",
    "import os\n",
    "if os.path.isfile(\"../models/brepnet.py\"):\n",
    "    os.chdir(\"../\")\n",
    "import tempfile\n",
    "import numpy as np\n",
    "    \n",
    "import utils.data_utils as data_utils\n",
    "\n",
    "# Imports from occwl\n",
    "from occwl.io import load_step\n",
    "from occwl.jupyter_viewer import JupyterViewer\n",
    "from occwl.entity_mapper import EntityMapper\n",
    "\n",
    "\n",
    "# Imports from BRepNet\n",
    "from pipeline.extract_brepnet_data_from_step import extract_brepnet_data_from_step\n",
    "import utils.scale_utils as scale_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1a0a88eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here is the path to some example step files for us to convert\n",
    "step_folder = Path(\"./example_files/step_examples\")\n",
    "temp_folder = Path(tempfile.gettempdir())\n",
    "working = temp_folder / \"brepnet_test_working_dir\"\n",
    "if not working.exists():\n",
    "    working.mkdir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "acbfe0aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 25/25 [02:51<00:00,  6.85s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed pipeline/extract_feature_data_from_step.py\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Here we generate npz data from step files\n",
    "extract_brepnet_data_from_step(step_folder, working)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e405436b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of step and npz files 25\n",
      "Viewing file\n",
      "npz data /tmp/brepnet_test_working_dir/30419_d55a0a22_2.npz\n",
      "STEP data example_files/step_examples/30419_d55a0a22_2.stp\n",
      "Loaded 1 solids\n"
     ]
    }
   ],
   "source": [
    "all_npz_files = [ f for f in working.glob(\"*.npz\") ]\n",
    "npz_files = []\n",
    "step_files = []\n",
    "for npz_file in all_npz_files:\n",
    "    step_file = step_folder / (npz_file.stem + \".stp\")\n",
    "    if step_file.exists():\n",
    "        npz_files.append(npz_file)\n",
    "        step_files.append(step_file)\n",
    "    \n",
    "print(f\"Number of step and npz files {len(npz_files)}\")\n",
    "\n",
    "file_index = 4\n",
    "npz_file = npz_files[file_index]\n",
    "step_file = step_folder / (npz_file.stem + \".stp\")\n",
    "print(\"Viewing file\") \n",
    "print(f\"npz data {npz_file}\")\n",
    "print(f\"STEP data {step_file}\")\n",
    "\n",
    "# Load the solid\n",
    "solids = load_step(step_file)\n",
    "print(f\"Loaded {len(solids)} solids\")\n",
    "solid = solids[0]\n",
    "\n",
    "# Scale to [-1,1]^3 box\n",
    "solid = scale_utils.scale_solid_to_unit_box(solid)\n",
    "\n",
    "# Load the npz data\n",
    "data = data_utils.load_npz_data(npz_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "33dd263f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb662b8ce1814f89b5e74a93a58e4fc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viewer = JupyterViewer()\n",
    "for face in solid.faces():\n",
    "    viewer.display(face, render_edges=True)\n",
    "viewer.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "15f5dc08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The faces you selected were [101]\n",
      "Shape of face grids (102, 7, 10, 10)\n",
      "Shape of selected face grids (1, 7, 10, 10)\n"
     ]
    }
   ],
   "source": [
    "if len(viewer.selected_faces()) <= 0:\n",
    "    print(\"Please select some faces to view this demo\")\n",
    "else:\n",
    "    entity_mapper = EntityMapper(solid)\n",
    "    selected_face_indices = viewer.selected_face_indices(entity_mapper)\n",
    "    print(f\"The faces you selected were {selected_face_indices}\")\n",
    "\n",
    "    # Get the point grids from the data\n",
    "    face_grids = data[\"face_point_grids\"]\n",
    "    print(f\"Shape of face grids {face_grids.shape}\")\n",
    "    selected_face_grids = face_grids[selected_face_indices]\n",
    "    print(f\"Shape of selected face grids {selected_face_grids.shape}\")\n",
    "    selected_face_grid_points = selected_face_grids[:,:3,:,:]\n",
    "    points = np.transpose(np.reshape(selected_face_grid_points, (3, -1)))\n",
    "    viewer.display_points(points)\n",
    "    selected_face_grid_normals  = selected_face_grids[:,3:6,:,:]\n",
    "    normals = np.transpose(np.reshape(selected_face_grid_normals, (3, -1)))\n",
    "    viewer.display_unit_vectors(points, normals)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72df579e",
   "metadata": {},
   "source": [
    "# View edge reverse flag as color\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "421ceab5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fc2c2660b514cd0a337271fe23fec6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "edge_color_viewer = JupyterViewer()\n",
    "edge_color_viewer.display(solid, render_edges=False)\n",
    "for edge in solid.edges():\n",
    "    if edge.reversed():\n",
    "        color = \"red\"\n",
    "    else:\n",
    "        color = \"blue\"\n",
    "    edge_color_viewer.display(edge, edge_color=color)\n",
    "edge_color_viewer.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ec2cbf",
   "metadata": {},
   "source": [
    "# View edge reverse flag as direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "05ca45e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfd9bcb67c7e4939955229b3f2a6f664",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def display_dir_across_edge(viewer, coedge_grid):\n",
    "    points = np.transpose(coedge_grid[:3, :])\n",
    "    tangents = np.transpose(coedge_grid[3:6, :])\n",
    "    left_normals = np.transpose(coedge_grid[6:9, :])\n",
    "    across_edge_dir = np.cross(tangents, left_normals)\n",
    "    viewer.display_points(points)\n",
    "    viewer.display_unit_vectors(points, across_edge_dir, line_color=\"blue\")\n",
    "\n",
    "    \n",
    "edge_dir_viewer = JupyterViewer()\n",
    "edge_dir_viewer.display(solid, render_edges=True)\n",
    "entity_mapper = EntityMapper(solid)\n",
    "coedge_grids = data[\"coedge_point_grids\"]\n",
    "for edge in solid.edges():\n",
    "    if edge.reversed():\n",
    "        coedge = edge.reversed_edge()\n",
    "    else:\n",
    "        coedge = edge\n",
    "    coedge_index = entity_mapper.oriented_edge_index(coedge)\n",
    "    coedge_grid = coedge_grids[coedge_index]\n",
    "    display_dir_across_edge(edge_dir_viewer, coedge_grid)\n",
    "    \n",
    "edge_dir_viewer.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fc63bb7",
   "metadata": {},
   "source": [
    "# View the coedge grids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1822f17a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad121aaa549d40328823f1efd7b5b336",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "coedge_viewer = JupyterViewer()\n",
    "for face in solid.faces():\n",
    "    coedge_viewer.display(face, render_edges=True)\n",
    "coedge_viewer.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb9be15d",
   "metadata": {},
   "source": [
    "For this demo you need to select two adjacent faces and then we can visualize the coedge grid for the common edge.  The coedge is on the side of the first face."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "47dac422",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 selected faces\n",
      "Face 0: 100\n",
      "Face 1: 36\n",
      "Found edges which shares faces\n",
      "Coedge index 550\n"
     ]
    }
   ],
   "source": [
    "def display_coedge_grid(viewer, grid):\n",
    "    points = np.transpose(grid[:3, :])\n",
    "    tangents = np.transpose(grid[3:6, :])\n",
    "    left_normals = np.transpose(grid[6:9, :])\n",
    "    right_normals = np.transpose(grid[9:12, :])\n",
    "    viewer.display_points(points)\n",
    "    viewer.display_unit_vectors(points, tangents, line_color=\"red\")\n",
    "    viewer.display_unit_vectors(points, left_normals, line_color=\"blue\")\n",
    "    viewer.display_unit_vectors(points, right_normals, line_color=\"green\")\n",
    "    \n",
    "def find_selected_coedge(viewer, entity_mapper):\n",
    "    selected_faces = viewer.selected_faces()\n",
    "    for index, selected_face in enumerate(selected_faces):\n",
    "        selected_face_index = entity_mapper.face_index(selected_face)\n",
    "        print(f\"Face {index}: {selected_face_index}\")\n",
    "        \n",
    "    selected_faces_set = set(selected_faces)\n",
    "       \n",
    "    # Find the coedge to the left of the face 0 and right of face 1\n",
    "    target_coedge = None\n",
    "    for coedge in solid.edges():\n",
    "        ffe = set(solid.faces_from_edge(coedge))\n",
    "        if ffe == selected_faces_set:\n",
    "            print(\"Found edges which shares faces\")\n",
    "            mate = coedge.reversed_edge()\n",
    "            if len(selected_faces) == 2:\n",
    "                if selected_faces[0].is_left_of(coedge) and selected_faces[1].is_left_of(mate):\n",
    "                    target_coedge = coedge\n",
    "                elif selected_faces[0].is_left_of(mate) and selected_faces[1].is_left_of(coedge):\n",
    "                    target_coedge = mate\n",
    "            elif len(selected_faces) == 1:\n",
    "                assert False, \"Need to implement\"\n",
    "    return target_coedge\n",
    "    \n",
    "num_selected_faces = len(coedge_viewer.selected_faces())\n",
    "print(f\"Found {num_selected_faces} selected faces\")\n",
    "if num_selected_faces > 0:\n",
    "    entity_mapper = EntityMapper(solid)\n",
    "    target_coedge = find_selected_coedge(coedge_viewer, entity_mapper)\n",
    "            \n",
    "    if target_coedge is not None:\n",
    "        coedge_index = entity_mapper.oriented_edge_index(target_coedge)\n",
    "        print(f\"Coedge index {coedge_index}\")\n",
    "        coedge_grids = data[\"coedge_point_grids\"]\n",
    "        selected_coedge_grids = coedge_grids[coedge_index]\n",
    "        display_coedge_grid(coedge_viewer, selected_coedge_grids)\n",
    "    else:\n",
    "        print(\"Found no coedge\")\n",
    "else:\n",
    "    print(\"Please select one face and one edge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85296550",
   "metadata": {},
   "source": [
    "# View the coedge LCS matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d29aa8fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa03c1b50311459b970f3e2cc29328d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "lcs_viewer = JupyterViewer()\n",
    "for face in solid.faces():\n",
    "    lcs_viewer.display(face, render_edges=True)\n",
    "lcs_viewer.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "065ad546",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 selected faces\n",
      "Face 0: 48\n",
      "Face 1: 100\n",
      "Found edges which shares faces\n",
      "Coedge index 258\n"
     ]
    }
   ],
   "source": [
    "def display_coedge_lcs(viewer, selected_coedge_lcs):\n",
    "    origin = np.expand_dims(selected_coedge_lcs[:, 3][:3], axis=0)\n",
    "    u_vec = np.expand_dims(selected_coedge_lcs[:, 0][:3], axis=0)\n",
    "    v_vec = np.expand_dims(selected_coedge_lcs[:, 1][:3], axis=0)\n",
    "    w_vec = np.expand_dims(selected_coedge_lcs[:, 2][:3], axis=0)\n",
    "    length = 0.02\n",
    "    u_end = origin + u_vec*length\n",
    "    v_end = origin + v_vec*length\n",
    "    w_end = origin + w_vec*length\n",
    "    viewer.display_points(origin)\n",
    "    viewer.display_lines(origin, u_end, line_width=2, line_color=\"red\")\n",
    "    viewer.display_lines(origin, v_end, line_width=2, line_color=\"green\")\n",
    "    viewer.display_lines(origin, w_end, line_width=2, line_color=\"blue\")\n",
    "\n",
    "num_selected_faces = len(lcs_viewer.selected_faces())\n",
    "print(f\"Found {num_selected_faces} selected faces\")\n",
    "if num_selected_faces > 0:\n",
    "    entity_mapper = EntityMapper(solid)\n",
    "    target_coedge = find_selected_coedge(lcs_viewer, entity_mapper)\n",
    "            \n",
    "    if target_coedge is not None:\n",
    "        coedge_index = entity_mapper.oriented_edge_index(target_coedge)\n",
    "        print(f\"Coedge index {coedge_index}\")\n",
    "        coedge_lcs = data[\"coedge_lcs\"]\n",
    "        selected_coedge_lcs = coedge_lcs[coedge_index]\n",
    "        display_coedge_lcs(lcs_viewer, selected_coedge_lcs)\n",
    "    else:\n",
    "        print(\"Found no coedge\")\n",
    "else:\n",
    "    print(\"Please select one face and one edge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b3aa400",
   "metadata": {},
   "source": [
    "# View edge grids which UV-Net sees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c1450b45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54e0f13d721a4e79bf4fb2e0e4f7fa0c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "uv_net_edge_viewer = JupyterViewer()\n",
    "uv_net_edge_viewer.display(solid, render_edges=True)\n",
    "entity_mapper = EntityMapper(solid)\n",
    "coedge_grids = data[\"coedge_point_grids\"]\n",
    "for edge in solid.edges():\n",
    "    coedge_index = entity_mapper.oriented_edge_index(edge)\n",
    "    coedge_grid = coedge_grids[coedge_index]\n",
    "    display_coedge_grid(uv_net_edge_viewer, coedge_grid)\n",
    "uv_net_edge_viewer.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15a04181",
   "metadata": {},
   "source": [
    "# View edge tangent directions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "90ca6f7f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48cdad7e2f5741a0ad2014bc60ec2866",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(HBox(children=(Checkbox(value=True, description='Axes', layout=Layout(height='au…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def display_edge_tangents(viewer, grid):\n",
    "    points = np.transpose(grid[:3, [0, 2, 4, 6, 8]])\n",
    "    \n",
    "    view_dir = np.array([1,1,-1])\n",
    "    \n",
    "    # Find a good length for the arrows\n",
    "    mins = np.min(points, axis=0)\n",
    "    maxs = np.max(points, axis=0)\n",
    "    diag = maxs - mins\n",
    "    longest = np.max(diag)\n",
    "    arrow_length = longest/8   \n",
    "    arrow_head = arrow_length/4\n",
    "    \n",
    "    points += view_dir*arrow_head\n",
    "    \n",
    "    tangents = np.transpose(grid[3:6, [0, 2, 4, 6, 8]])\n",
    "    arrow_ends = points + arrow_length*tangents\n",
    "    arrow_dir_cross = view_dir/np.linalg.norm(view_dir)\n",
    "    arrow_offset1 = np.cross(tangents, arrow_dir_cross)\n",
    "    arrow_offset2 = -arrow_offset1\n",
    "    arrow1 = arrow_ends - (tangents+arrow_offset1)*arrow_head\n",
    "    arrow2 = arrow_ends - (tangents+arrow_offset2)*arrow_head\n",
    "    line_width = 2\n",
    "    viewer.display_lines(points, arrow_ends, line_color=\"blue\", line_width=line_width)\n",
    "    viewer.display_lines(arrow1, arrow_ends, line_color=\"blue\", line_width=line_width)\n",
    "    viewer.display_lines(arrow2, arrow_ends, line_color=\"blue\", line_width=line_width)\n",
    "    \n",
    "edge_tangent_viewer = JupyterViewer()\n",
    "edge_tangent_viewer.display(solid, render_edges=True)\n",
    "entity_mapper = EntityMapper(solid)\n",
    "coedge_grids = data[\"coedge_point_grids\"]\n",
    "for edge in solid.edges():\n",
    "    if edge.reversed():\n",
    "        # If the edge had the reverse flag set then the code in \n",
    "        # the EdgeDataExtractor already reversed the tangent directions.\n",
    "        # To get the same tangent directions as the edges 3d curve\n",
    "        # we need to get the coedge grid from the mate\n",
    "        coedge = edge.reversed_edge()\n",
    "    else:\n",
    "        coedge = edge\n",
    "    coedge_index = entity_mapper.oriented_edge_index(coedge)\n",
    "    coedge_grid = coedge_grids[coedge_index]\n",
    "    display_edge_tangents(edge_tangent_viewer, coedge_grid)\n",
    "edge_tangent_viewer.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}