2.1 - Loading-Shapefiles-To-VTK.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shapefile\n",
    "import numpy as np\n",
    "\n",
    "import vtk\n",
    "import vista\n",
    "\n",
    "# sklearn's KDTree is fast: use it if available\n",
    "# Scipy also also a good KDTree\n",
    "from sklearn.neighbors import KDTree as Tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define conveinance functions for converting shape files to VTK objects\n",
    "\n",
    "def _fix_to_topography(topo_points, points_to_update, static=20.0):\n",
    "    \"\"\"Update the z component of points to force them to lie on a topo surface\"\"\"\n",
    "    tree = Tree(topo_points)\n",
    "    ind = tree.query(points_to_update, k=1)[1].ravel()\n",
    "    # Now update the elevation to be on the topo surface\n",
    "    # Also shift it so its always just above the surface and not in the surface\n",
    "    points_to_update[:,2] = topo_points[:,2][ind] + static\n",
    "    return points_to_update\n",
    "\n",
    "def _makeLineCell(idx0, idx1):\n",
    "    \"\"\"Create a vtkLine cell\"\"\"\n",
    "    aLine = vtk.vtkLine()\n",
    "    aLine.GetPointIds().SetId(0, idx0)\n",
    "    aLine.GetPointIds().SetId(1, idx1)\n",
    "    return aLine\n",
    "\n",
    "def polygon_to_vtk(polygon, topo_points=None):\n",
    "    \"\"\"Converts a polygon shape to a vista.PolyData object.\n",
    "    This assumes the points are ordered.\n",
    "    \"\"\"\n",
    "    pts = np.array(polygon.points)\n",
    "    pts = np.c_[pts, np.zeros(pts.shape[0])]\n",
    "    if topo_points is not None:\n",
    "        pts = _fix_to_topography(topo_points, pts)\n",
    "        \n",
    "    cells = vtk.vtkCellArray()\n",
    "    for i in range(pts.shape[0]-1):\n",
    "        cell = _makeLineCell(i, i+1)\n",
    "        cells.InsertNextCell(cell)\n",
    "    \n",
    "    # Add in last connection to make complete polygon\n",
    "    cell = _makeLineCell(i, 0)\n",
    "    cells.InsertNextCell(cell)\n",
    "    \n",
    "    # Build the output\n",
    "    pdo = vtk.vtkPolyData()\n",
    "    pdo.SetPoints(vista.vtk_points(pts))\n",
    "    pdo.SetLines(cells)\n",
    "    return vista.wrap(pdo)\n",
    "\n",
    "VTK_CONVERTERS = {\n",
    "    shapefile.POLYGON: polygon_to_vtk,\n",
    "}\n",
    "\n",
    "def read_shape_file_to_vtk(filename, topo_points=None):\n",
    "    \"\"\"Read all the features of a shapefile into vista objects.\n",
    "    Use the topo_points argument to fill the Z component of 2D points\n",
    "    \"\"\"\n",
    "    shp = shapefile.Reader(filename)\n",
    "    output = vista.MultiBlock()\n",
    "    for i, feature in enumerate(shp.shapeRecords()):\n",
    "        shape = feature.shape\n",
    "        print(feature.record[1])\n",
    "        try:\n",
    "            try:\n",
    "                output[i, feature.record[1]] = VTK_CONVERTERS[shape.shapeType](shape, topo_points)\n",
    "            except:\n",
    "                output[i] = VTK_CONVERTERS[shape.shapeType](shape, topo_points)\n",
    "            return output\n",
    "        except KeyError:\n",
    "            raise RuntimeError('Shape type ({}) unknown'.format(shape.shapeType))\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = './data/shapefiles/outline'\n",
    "shapes = read_shape_file_to_vtk(filename)\n",
    "shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:dev]",
   "language": "python",
   "name": "dev"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}