GenerateData/4_make_auxiliary_dat_file.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NOTE\n",
    "\n",
    "Please note that this script uses python3, which should still work in the python2 version, but be careful to adjust the usage of the library to match compatibility.\n",
    "\n",
    "The original version of Pixel2Mesh's `.dat` file has pkl[5] and pkl[6], but they are not used by other code, so only this part is padded to 0 here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import networkx as nx\n",
    "import scipy.sparse as sp\n",
    "import sys\n",
    "import os\n",
    "import pickle\n",
    "import trimesh\n",
    "\n",
    "from IPython.display import Image\n",
    "from scipy.sparse.linalg.eigen.arpack import eigsh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_obj(fn, no_normal=False):\n",
    "    fin = open(fn, 'r')\n",
    "    lines = [line.rstrip() for line in fin]\n",
    "    fin.close()\n",
    "\n",
    "    vertices = []; normals = []; faces = [];\n",
    "    for line in lines:\n",
    "        if line.startswith('v '):\n",
    "            vertices.append(np.float32(line.split()[1:4]))\n",
    "        elif line.startswith('vn '):\n",
    "            normals.append(np.float32(line.split()[1:4]))\n",
    "        elif line.startswith('f '):\n",
    "            faces.append(np.int32([item.split('/')[0] for item in line.split()[1:4]]))\n",
    "\n",
    "    mesh = dict()\n",
    "    mesh['faces'] = np.vstack(faces)\n",
    "    mesh['vertices'] = np.vstack(vertices)\n",
    "\n",
    "    if (not no_normal) and (len(normals) > 0):\n",
    "        assert len(normals) == len(vertices), 'ERROR: #vertices != #normals'\n",
    "        mesh['normals'] = np.vstack(normals)\n",
    "\n",
    "    return mesh\n",
    "\n",
    "def sparse_to_tuple(sparse_mx):\n",
    "    \"\"\"Convert sparse matrix to tuple representation.\"\"\"\n",
    "    def to_tuple(mx):\n",
    "        if not sp.isspmatrix_coo(mx):\n",
    "            mx = mx.tocoo()\n",
    "        coords = np.vstack((mx.row, mx.col)).transpose()\n",
    "        values = mx.data\n",
    "        shape = mx.shape\n",
    "        return coords, values, shape\n",
    "\n",
    "    if isinstance(sparse_mx, list):\n",
    "        for i in range(len(sparse_mx)):\n",
    "            sparse_mx[i] = to_tuple(sparse_mx[i])\n",
    "    else:\n",
    "        sparse_mx = to_tuple(sparse_mx)\n",
    "\n",
    "    return sparse_mx\n",
    "\n",
    "\n",
    "def normalize_adj(adj):\n",
    "    \"\"\"Symmetrically normalize adjacency matrix.\"\"\"\n",
    "    adj = sp.coo_matrix(adj)\n",
    "    rowsum = np.array(adj.sum(1))\n",
    "    d_inv_sqrt = np.power(rowsum, -0.5).flatten()\n",
    "    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.\n",
    "    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)\n",
    "    return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()\n",
    "\n",
    "\n",
    "def preprocess_adj(adj):\n",
    "    \"\"\"Preprocessing of adjacency matrix for simple GCN model and conversion to tuple representation.\"\"\"\n",
    "    adj_normalized = normalize_adj(adj + sp.eye(adj.shape[0]))\n",
    "    return sparse_to_tuple(adj_normalized)\n",
    "\n",
    "\n",
    "def construct_feed_dict(features, support, labels, labels_mask, placeholders):\n",
    "    \"\"\"Construct feed dictionary.\"\"\"\n",
    "    feed_dict = dict()\n",
    "    feed_dict.update({placeholders['labels']: labels})\n",
    "    feed_dict.update({placeholders['labels_mask']: labels_mask})\n",
    "    feed_dict.update({placeholders['features']: features})\n",
    "    feed_dict.update({placeholders['support'][i]: support[i] for i in range(len(support))})\n",
    "    feed_dict.update({placeholders['num_features_nonzero']: features[1].shape})\n",
    "    return feed_dict\n",
    "\n",
    "\n",
    "def chebyshev_polynomials(adj, k):\n",
    "    \"\"\"Calculate Chebyshev polynomials up to order k. Return a list of sparse matrices (tuple representation).\"\"\"\n",
    "    print(\"Calculating Chebyshev polynomials up to order {}...\".format(k))\n",
    "\n",
    "    adj_normalized = normalize_adj(adj)\n",
    "    laplacian = sp.eye(adj.shape[0]) - adj_normalized\n",
    "    largest_eigval, _ = eigsh(laplacian, 1, which='LM')\n",
    "    scaled_laplacian = (2. / largest_eigval[0]) * laplacian - sp.eye(adj.shape[0])\n",
    "\n",
    "    t_k = list()\n",
    "    t_k.append(sp.eye(adj.shape[0]))\n",
    "    t_k.append(scaled_laplacian)\n",
    "\n",
    "    def chebyshev_recurrence(t_k_minus_one, t_k_minus_two, scaled_lap):\n",
    "        s_lap = sp.csr_matrix(scaled_lap, copy=True)\n",
    "        return 2 * s_lap.dot(t_k_minus_one) - t_k_minus_two\n",
    "\n",
    "    for i in range(2, k+1):\n",
    "        t_k.append(chebyshev_recurrence(t_k[-1], t_k[-2], scaled_laplacian))\n",
    "\n",
    "    return sparse_to_tuple(t_k)\n",
    "\n",
    "\n",
    "def dense_cheb(adj, k):\n",
    "    \"\"\"Calculate Chebyshev polynomials up to order k. Return a list of sparse matrices (tuple representation).\"\"\"\n",
    "    print(\"Calculating Chebyshev polynomials up to order {}...\".format(k))\n",
    "\n",
    "    adj_normalized = normalize_adj(adj)\n",
    "    laplacian = sp.eye(adj.shape[0]) - adj_normalized\n",
    "    largest_eigval, _ = eigsh(laplacian, 1, which='LM')\n",
    "    scaled_laplacian = (2. / largest_eigval[0]) * laplacian - sp.eye(adj.shape[0])\n",
    "\n",
    "    t_k = list()\n",
    "    t_k.append(sp.eye(adj.shape[0]))\n",
    "    t_k.append(scaled_laplacian)\n",
    "\n",
    "    def chebyshev_recurrence(t_k_minus_one, t_k_minus_two, scaled_lap):\n",
    "        s_lap = sp.csr_matrix(scaled_lap, copy=True)\n",
    "        return 2 * s_lap.dot(t_k_minus_one) - t_k_minus_two\n",
    "\n",
    "    for i in range(2, k+1):\n",
    "        t_k.append(chebyshev_recurrence(t_k[-1], t_k[-2], scaled_laplacian))\n",
    "\n",
    "    return t_k\n",
    "\n",
    "def unpool_face(old_faces, old_unique_edges, old_vertices):\n",
    "    old_faces = np.array(old_faces)\n",
    "    N = old_vertices.shape[0]\n",
    "    mid_table = np.zeros((N,N), dtype=np.int32)\n",
    "    new_edges = []\n",
    "    new_faces = []\n",
    "    for i, u in enumerate(old_unique_edges):\n",
    "        mid_table[u[0], u[1]] = N+i\n",
    "        mid_table[u[1], u[0]] = N+i\n",
    "        new_edges.append([u[0], N+i])\n",
    "        new_edges.append([N+i, u[1]])\n",
    "    \n",
    "    for i, f in enumerate(old_faces):\n",
    "        f = np.sort(f)\n",
    "        mid1 = mid_table[f[0], f[1]]\n",
    "        mid2 = mid_table[f[0], f[2]]\n",
    "        mid3 = mid_table[f[1], f[2]]\n",
    "        \n",
    "        new_faces.append([f[0], mid1, mid2])\n",
    "        new_faces.append([f[1], mid1, mid3])\n",
    "        new_faces.append([f[2], mid2, mid3])\n",
    "        new_faces.append([mid1, mid2, mid3])\n",
    "        \n",
    "        new_edges.append([mid1, mid2])\n",
    "        new_edges.append([mid2, mid3])\n",
    "        new_edges.append([mid3, mid1])\n",
    "    \n",
    "    new_faces = np.array(new_faces, dtype=np.int32)\n",
    "    new_edges = np.array(new_edges, dtype=np.int32)\n",
    "    return new_edges, new_faces\n",
    "\n",
    "\n",
    "def write_obj(path, vertices, faces):\n",
    "    with open(path, 'w') as o:\n",
    "        for v in vertices:\n",
    "            o.write('v {} {} {}\\n'.format(v[0], v[1], v[2]))\n",
    "        for f in faces:\n",
    "            o.write('f {} {} {}\\n'.format(f[0]+1, f[1]+1, f[2]+1))\n",
    "            \n",
    "\n",
    "def cal_lap_index(mesh_neighbor):\n",
    "    lap_index = np.zeros([mesh_neighbor.shape[0], 2 + 8]).astype(np.int32)\n",
    "    for i, j in enumerate(mesh_neighbor):\n",
    "        lenj = len(j)\n",
    "        lap_index[i][0:lenj] = j\n",
    "        lap_index[i][lenj:-2] = -1\n",
    "        lap_index[i][-2] = i\n",
    "        lap_index[i][-1] = lenj\n",
    "    return lap_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pkl = pickle.load(open('../Data/ellipsoid/info_ellipsoid.dat', 'rb'), encoding='bytes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "info = {}\n",
    "info['coords'] = None\n",
    "info['support'] = {'stage1':None,'stage2':None,'stage3':None, 'stage4':None}\n",
    "info['unpool_idx'] = {'stage1_2':None,'stage2_3':None, 'stage3_4':None}\n",
    "info['lap_idx'] = {'stage1':None,'stage2':None,'stage3':None,'stage4':None}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simply load obj file created by Meshlab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_mesh = load_obj('./init_obj/init1.obj',no_normal=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reload mesh using trimesh to get adjacent matrix, set `process=Flase` to preserve mesh vertices order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = trimesh.Trimesh(vertices=raw_mesh['vertices'], faces=(raw_mesh['faces']-1), process=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.all(raw_mesh['faces'] == mesh.faces+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords_1 = np.array(mesh.vertices, dtype=np.float32)\n",
    "info['coords'] = coords_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 1 auxiliary matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating Chebyshev polynomials up to order 1...\n"
     ]
    }
   ],
   "source": [
    "adj_1 = nx.adjacency_matrix(mesh.vertex_adjacency_graph, nodelist=range(len(coords_1)))\n",
    "cheb_1 = chebyshev_polynomials(adj_1,1)\n",
    "info['support']['stage1'] = cheb_1\n",
    "\n",
    "edges_1 = mesh.edges_unique\n",
    "edges_1 = edges_1[edges_1[:,1].argsort(kind='mergesort')]\n",
    "edges_1 = edges_1[edges_1[:,0].argsort(kind='mergesort')]\n",
    "info['unpool_idx']['stage1_2'] = edges_1\n",
    "\n",
    "lap_1 = cal_lap_index(mesh.vertex_neighbors)\n",
    "info['lap_idx']['stage1'] = lap_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 2 auxiliary matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating Chebyshev polynomials up to order 1...\n"
     ]
    }
   ],
   "source": [
    "faces_1 = np.array(mesh.faces)\n",
    "\n",
    "edges_2, faces_2 = unpool_face(faces_1, edges_1, coords_1)\n",
    "\n",
    "tmp_1_2 = 0.5*(coords_1[info['unpool_idx']['stage1_2'][:,0]] + coords_1[info['unpool_idx']['stage1_2'][:,1]])\n",
    "coords_2 = np.vstack([coords_1, tmp_1_2])\n",
    "\n",
    "mesh2 = trimesh.Trimesh(vertices=coords_2, faces=faces_2, process=False)\n",
    "\n",
    "adj_2 = nx.adjacency_matrix(mesh2.vertex_adjacency_graph, nodelist=range(len(coords_2)))\n",
    "cheb_2 = chebyshev_polynomials(adj_2,1)\n",
    "info['support']['stage2'] = cheb_2\n",
    "\n",
    "edges_2 = edges_2[edges_2[:,1].argsort(kind='mergesort')]\n",
    "edges_2 = edges_2[edges_2[:,0].argsort(kind='mergesort')]\n",
    "info['unpool_idx']['stage2_3'] = edges_2\n",
    "\n",
    "lap_2 = cal_lap_index(mesh2.vertex_neighbors)\n",
    "info['lap_idx']['stage2'] = lap_2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save init2.obj, you can only save faces to get face2.obj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_obj('./init_obj/init2.obj', coords_2, faces_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 3 auxiliary matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating Chebyshev polynomials up to order 1...\n"
     ]
    }
   ],
   "source": [
    "edges_3, faces_3 = unpool_face(faces_2, edges_2, coords_2)\n",
    "\n",
    "tmp_2_3 = 0.5*(coords_2[info['unpool_idx']['stage2_3'][:,0]] + coords_2[info['unpool_idx']['stage2_3'][:,1]])\n",
    "coords_3 = np.vstack([coords_2, tmp_2_3])\n",
    "\n",
    "mesh3 = trimesh.Trimesh(vertices=coords_3, faces=faces_3, process=False)\n",
    "\n",
    "adj_3 = nx.adjacency_matrix(mesh3.vertex_adjacency_graph, nodelist=range(len(coords_3)))\n",
    "cheb_3 = chebyshev_polynomials(adj_3,1)\n",
    "info['support']['stage3'] = cheb_3\n",
    "\n",
    "edges_3 = edges_3[edges_3[:,1].argsort(kind='mergesort')]\n",
    "edges_3 = edges_3[edges_3[:,0].argsort(kind='mergesort')]\n",
    "info['unpool_idx']['stage3_4'] = edges_3\n",
    "\n",
    "lap_3 = cal_lap_index(mesh3.vertex_neighbors)\n",
    "info['lap_idx']['stage3'] = lap_3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save init3.obj, you can only save faces to get face3.obj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_obj('./init_obj/init3.obj', coords_3, faces_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stage 4 auxiliary matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating Chebyshev polynomials up to order 1...\n"
     ]
    }
   ],
   "source": [
    "edges_4, faces_4 = unpool_face(faces_3, edges_3, coords_3)\n",
    "\n",
    "tmp_3_4 = 0.5*(coords_3[info['unpool_idx']['stage3_4'][:,0]] + coords_3[info['unpool_idx']['stage3_4'][:,1]])\n",
    "coords_4 = np.vstack([coords_3, tmp_3_4])\n",
    "\n",
    "mesh4 = trimesh.Trimesh(vertices=coords_4, faces=faces_4, process=False)\n",
    "\n",
    "adj_4 = nx.adjacency_matrix(mesh4.vertex_adjacency_graph, nodelist=range(len(coords_4)))\n",
    "cheb_4 = chebyshev_polynomials(adj_4,1)\n",
    "info['support']['stage4'] = cheb_4\n",
    "\n",
    "edges_4 = edges_4[edges_4[:,1].argsort(kind='mergesort')]\n",
    "edges_4 = edges_4[edges_4[:,0].argsort(kind='mergesort')]\n",
    "info['unpool_idx']['stage4_5'] = edges_4\n",
    "\n",
    "lap_4 = cal_lap_index(mesh4.vertex_neighbors)\n",
    "info['lap_idx']['stage4'] = lap_4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_obj('./init_obj/init4.obj', coords_4, faces_4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dump .dat file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat = [info['coords'],\n",
    "       info['support']['stage1'],\n",
    "       info['support']['stage2'],\n",
    "       info['support']['stage3'],\n",
    "       info['support']['stage4'],\n",
    "       [info['unpool_idx']['stage1_2'],\n",
    "        info['unpool_idx']['stage2_3'],\n",
    "        info['unpool_idx']['stage3_4']\n",
    "       ],\n",
    "       [np.zeros((1,4), dtype=np.int32)]*4,\n",
    "       [np.zeros((1,4))]*4,\n",
    "       [info['lap_idx']['stage1'],\n",
    "        info['lap_idx']['stage2'],\n",
    "        info['lap_idx']['stage3'],\n",
    "        info['lap_idx']['stage4']\n",
    "       ],\n",
    "      ]\n",
    "pickle.dump(dat, open(\"./init_obj/pixel2mesh_aux_4stages.dat\",\"wb\"), protocol=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}