{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Matching Image Keypoints by Graph Matching Neural Networks\n\nThis example shows how to match image keypoints by neural network-based graph matching solvers.\nThese graph matching solvers are designed to match two individual graphs. The matched images\ncan be further passed to tackle downstream tasks.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Runzhong Wang <runzhong.wang@sjtu.edu.cn>\n#         Wenzheng Pan <pwz1121@sjtu.edu.cn>\n#\n# License: Mulan PSL v2 License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>The following solvers are based on matching two individual graphs, and are included in this example:\n\n    * :func:`~pygmtools.neural_solvers.pca_gm` (neural network solver)\n\n    * :func:`~pygmtools.neural_solvers.ipca_gm` (neural network solver)\n\n    * :func:`~pygmtools.neural_solvers.cie` (neural network solver)</p></div>\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import jittor as jt # jittor backend\nfrom jittor import Var, models, nn\nimport pygmtools as pygm\nimport matplotlib.pyplot as plt # for plotting\nfrom matplotlib.patches import ConnectionPatch # for plotting matching result\nimport scipy.io as sio # for loading .mat file\nimport scipy.spatial as spa # for Delaunay triangulation\nfrom sklearn.decomposition import PCA as PCAdimReduc\nimport itertools\nimport numpy as np\nfrom PIL import Image\npygm.BACKEND = 'jittor' # set default backend for pygmtools\njt.flags.use_cuda = jt.has_cuda"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Predicting Matching by Graph Matching Neural Networks\nIn this section we show how to do predictions (inference) by graph matching neural networks.\nLet's take PCA-GM (:func:`~pygmtools.neural_solvers.pca_gm`) as an example.\n\n### Load the images\nImages are from the Willow Object Class dataset (this dataset also available with the Benchmark of ``pygmtools``,\nsee :class:`~pygmtools.dataset.WillowObject`).\n\nThe images are resized to 256x256.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "obj_resize = (256, 256)\nimg1 = Image.open('../data/willow_duck_0001.png')\nimg2 = Image.open('../data/willow_duck_0002.png')\nkpts1 = jt.Var(sio.loadmat('../data/willow_duck_0001.mat')['pts_coord'])\nkpts2 = jt.Var(sio.loadmat('../data/willow_duck_0002.mat')['pts_coord'])\nkpts1[0] = kpts1[0] * obj_resize[0] / img1.size[0]\nkpts1[1] = kpts1[1] * obj_resize[1] / img1.size[1]\nkpts2[0] = kpts2[0] * obj_resize[0] / img2.size[0]\nkpts2[1] = kpts2[1] * obj_resize[1] / img2.size[1]\nimg1 = img1.resize(obj_resize, resample=Image.BILINEAR)\nimg2 = img2.resize(obj_resize, resample=Image.BILINEAR)\njittor_img1 = jt.Var(np.array(img1, dtype=np.float32) / 256).permute(2, 0, 1).unsqueeze(0) # shape: BxCxHxW\njittor_img2 = jt.Var(np.array(img2, dtype=np.float32) / 256).permute(2, 0, 1).unsqueeze(0) # shape: BxCxHxW"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualize the images and keypoints\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def plot_image_with_graph(img, kpt, A=None):\n    plt.imshow(img)\n    plt.scatter(kpt[0], kpt[1], c='w', edgecolors='k')\n    if A is not None:\n        for idx in jt.nonzero(A):\n            plt.plot((kpt[0, idx[0]], kpt[0, idx[1]]), (kpt[1, idx[0]], kpt[1, idx[1]]), 'k-')\n\nplt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title('Image 1')\nplot_image_with_graph(img1, kpts1)\nplt.subplot(1, 2, 2)\nplt.title('Image 2')\nplot_image_with_graph(img2, kpts2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Build the graphs\nGraph structures are built based on the geometric structure of the keypoint set. In this example,\nwe refer to [Delaunay triangulation](https://en.wikipedia.org/wiki/Delaunay_triangulation).\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def delaunay_triangulation(kpt):\n    d = spa.Delaunay(kpt.numpy().transpose())\n    A = jt.zeros((len(kpt[0]), len(kpt[0])))\n    for simplex in d.simplices:\n        for pair in itertools.permutations(simplex, 2):\n            A[pair] = 1\n    return A\n\nA1 = delaunay_triangulation(kpts1)\nA2 = delaunay_triangulation(kpts2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualize the graphs\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title('Image 1 with Graphs')\nplot_image_with_graph(img1, kpts1, A1)\nplt.subplot(1, 2, 2)\nplt.title('Image 2 with Graphs')\nplot_image_with_graph(img2, kpts2, A2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Extract node features via CNN\nDeep graph matching solvers can be fused with CNN feature extractors, to build an end-to-end learning pipeline.\n\nIn this example, let's adopt the deep graph solvers based on matching two individual graphs.\nThe image features are based on two intermediate layers from the VGG16 CNN model, following\nexisting deep graph matching papers (such as :func:`~pygmtools.neural_solvers.pca_gm`)\n\nLet's firstly fetch and download the VGG16 model:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "vgg16_cnn = models.vgg16_bn(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "List of layers of VGG16:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(vgg16_cnn.features)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define the CNN feature extractor, which outputs the features of ``layer (30)`` and\n``layer (37)``\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class CNNNet(jt.nn.Module):\n    def __init__(self, vgg16_module):\n        super(CNNNet, self).__init__()\n        # The naming of the layers follow ThinkMatch convention to load pretrained models.\n        self.node_layers = jt.nn.Sequential(*[_ for _ in list(vgg16_module.features)[:31]])\n        self.edge_layers = jt.nn.Sequential(*[_ for _ in list(vgg16_module.features)[31:38]])\n\n    def execute(self, inp_img):\n        feat_local = self.node_layers(inp_img)\n        feat_global = self.edge_layers(feat_local)\n        return feat_local, feat_global"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download pretrained CNN weights (from [ThinkMatch](https://github.com/Thinklab-SJTU/ThinkMatch)),\nload the weights and then extract the CNN features\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "cnn = CNNNet(vgg16_cnn)\npath = pygm.utils.download('vgg16_pca_voc_jittor.pt', 'https://drive.google.com/u/0/uc?export=download&confirm=Z-AR&id=1qLxjcVq7X3brylxRJvELCbtCzfuXQ24J')\ncnn.load_state_dict(jt.load(path))\n\nwith jt.no_grad():\n    feat1_local, feat1_global = cnn(jittor_img1)\n    feat2_local, feat2_global = cnn(jittor_img2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Normalize the features\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def local_response_norm(input: Var, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1.0) -> Var:\n    \"\"\"\n    jittor implementation of local_response_norm\n    \"\"\"\n    dim = input.ndim\n    assert dim >= 3\n\n    if input.numel() == 0:\n        return input\n\n    div = input.multiply(input).unsqueeze(1)\n    if dim == 3:\n        div = nn.pad(div, (0, 0, size // 2, (size - 1) // 2))\n        div = nn.avg_pool2d(div, (size, 1), stride=1).squeeze(1)\n    else:\n        sizes = input.size()\n        div = div.view(sizes[0], 1, sizes[1], sizes[2], -1)\n        div = nn.pad(div, (0, 0, 0, 0, size // 2, (size - 1) // 2))\n        div = nn.AvgPool3d((size, 1, 1), stride=1)(div).squeeze(1)\n        div = div.view(sizes)\n    div = div.multiply(alpha).add(k).pow(beta)\n    return input / div\n\n\ndef l2norm(node_feat):\n    return local_response_norm(\n        node_feat, node_feat.shape[1] * 2, alpha=node_feat.shape[1] * 2, beta=0.5, k=0)\n\nfeat1_local = l2norm(feat1_local)\nfeat1_global = l2norm(feat1_global)\nfeat2_local = l2norm(feat2_local)\nfeat2_global = l2norm(feat2_global)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Up-sample the features to the original image size and concatenate\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "feat1_local_upsample = jt.nn.interpolate(feat1_local, obj_resize, mode='bilinear')\nfeat1_global_upsample = jt.nn.interpolate(feat1_global, obj_resize, mode='bilinear')\nfeat2_local_upsample = jt.nn.interpolate(feat2_local, obj_resize, mode='bilinear')\nfeat2_global_upsample = jt.nn.interpolate(feat2_global, obj_resize, mode='bilinear')\nfeat1_upsample = jt.concat((feat1_local_upsample, feat1_global_upsample), dim=1)\nfeat2_upsample = jt.concat((feat2_local_upsample, feat2_global_upsample), dim=1)\nnum_features = feat1_upsample.shape[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualize the extracted CNN feature (dimensionality reduction via principle component analysis)\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pca_dim_reduc = PCAdimReduc(n_components=3, whiten=True)\nfeat_dim_reduc = pca_dim_reduc.fit_transform(\n    np.concatenate((\n        feat1_upsample.permute(0, 2, 3, 1).reshape(-1, num_features).numpy(),\n        feat2_upsample.permute(0, 2, 3, 1).reshape(-1, num_features).numpy()\n    ), axis=0)\n)\nfeat_dim_reduc = feat_dim_reduc / np.max(np.abs(feat_dim_reduc), axis=0, keepdims=True) / 2 + 0.5\nfeat1_dim_reduc = feat_dim_reduc[:obj_resize[0] * obj_resize[1], :]\nfeat2_dim_reduc = feat_dim_reduc[obj_resize[0] * obj_resize[1]:, :]\n\nplt.figure(figsize=(8, 4))\nplt.subplot(1, 2, 1)\nplt.title('Image 1 with CNN features')\nplot_image_with_graph(img1, kpts1, A1)\nplt.imshow(feat1_dim_reduc.reshape(obj_resize[0], obj_resize[1], 3), alpha=0.5)\nplt.subplot(1, 2, 2)\nplt.title('Image 2 with CNN features')\nplot_image_with_graph(img2, kpts2, A2)\nplt.imshow(feat2_dim_reduc.reshape(obj_resize[0], obj_resize[1], 3), alpha=0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extract node features by nearest interpolation\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rounded_kpts1 = jt.round(kpts1).long()\nrounded_kpts2 = jt.round(kpts2).long()\nnode1 = feat1_upsample[0, :, rounded_kpts1[1], rounded_kpts1[0]].t() # shape: NxC\nnode2 = feat2_upsample[0, :, rounded_kpts2[1], rounded_kpts2[0]].t() # shape: NxC"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Call PCA-GM matching model\nSee :func:`~pygmtools.neural_solvers.pca_gm` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.pca_gm(node1, node2, A1, A2, pretrain='voc')\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by PCA-GM')\nax1 = plt.subplot(1, 2, 1)\nplot_image_with_graph(img1, kpts1, A1)\nax2 = plt.subplot(1, 2, 2)\nplot_image_with_graph(img2, kpts2, A2)\nidx, _ = jt.argmax(X, dim=1)\nfor i in range(X.shape[0]):\n    j = idx[i].item()\n    con = ConnectionPatch(xyA=kpts1[:, i], xyB=kpts2[:, j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"red\" if i != j else \"green\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Matching images with other neural networks\nThe above pipeline also works for other deep graph matching networks. Here we give examples of\n:func:`~pygmtoools.neural_solvers.ipca_gm` and :func:`~pygmtoools.neural_solvers.cie`.\n\n### Matching by IPCA-GM model\nSee :func:`~pygmtools.neural_solvers.ipca_gm` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "path = pygm.utils.download('vgg16_ipca_voc_jittor.pt', 'https://drive.google.com/u/0/uc?export=download&confirm=Z-AR&id=1f7KEl9ZFZwI26j6UId-fsdl8Y8QWPKZi')\ncnn.load_state_dict(jt.load(path))\n\nfeat1_local, feat1_global = cnn(jittor_img1)\nfeat2_local, feat2_global = cnn(jittor_img2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Normalize the features\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def l2norm(node_feat):\n    return local_response_norm(\n        node_feat, node_feat.shape[1] * 2, alpha=node_feat.shape[1] * 2, beta=0.5, k=0)\n\nfeat1_local = l2norm(feat1_local)\nfeat1_global = l2norm(feat1_global)\nfeat2_local = l2norm(feat2_local)\nfeat2_global = l2norm(feat2_global)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Up-sample the features to the original image size and concatenate\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "feat1_local_upsample = jt.nn.interpolate(feat1_local, obj_resize, mode='bilinear')\nfeat1_global_upsample = jt.nn.interpolate(feat1_global, obj_resize, mode='bilinear')\nfeat2_local_upsample = jt.nn.interpolate(feat2_local, obj_resize, mode='bilinear')\nfeat2_global_upsample = jt.nn.interpolate(feat2_global, obj_resize, mode='bilinear')\nfeat1_upsample = jt.concat((feat1_local_upsample, feat1_global_upsample), dim=1)\nfeat2_upsample = jt.concat((feat2_local_upsample, feat2_global_upsample), dim=1)\nnum_features = feat1_upsample.shape[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extract node features by nearest interpolation\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rounded_kpts1 = jt.round(kpts1).long()\nrounded_kpts2 = jt.round(kpts2).long()\nnode1 = feat1_upsample[0, :, rounded_kpts1[1], rounded_kpts1[0]].t() # shape: NxC\nnode2 = feat2_upsample[0, :, rounded_kpts2[1], rounded_kpts2[0]].t() # shape: NxC"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Build edge features as edge lengths\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kpts1_dis = (kpts1.unsqueeze(0) - kpts1.unsqueeze(1))\nkpts1_dis = jt.norm(kpts1_dis, p=2, dim=2).detach()\nkpts2_dis = (kpts2.unsqueeze(0) - kpts2.unsqueeze(1))\nkpts2_dis = jt.norm(kpts2_dis, p=2, dim=2).detach()\n\nQ1 = jt.exp(-kpts1_dis / obj_resize[0])\nQ2 = jt.exp(-kpts2_dis / obj_resize[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Matching by IPCA-GM model\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.ipca_gm(node1, node2, A1, A2, pretrain='voc')\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by IPCA-GM')\nax1 = plt.subplot(1, 2, 1)\nplot_image_with_graph(img1, kpts1, A1)\nax2 = plt.subplot(1, 2, 2)\nplot_image_with_graph(img2, kpts2, A2)\nidx, _ = jt.argmax(X, dim=1)\nfor i in range(X.shape[0]):\n    j = idx[i].item()\n    con = ConnectionPatch(xyA=kpts1[:, i], xyB=kpts2[:, j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"red\" if i != j else \"green\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Matching by CIE model\nSee :func:`~pygmtools.neural_solvers.cie` for the API reference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "path = pygm.utils.download('vgg16_cie_voc_jittor.pt', 'https://drive.google.com/u/0/uc?export=download&confirm=Z-AR&id=1wDbA-8sK4BNhA48z2c-Gtdd4AarRxfqT')\ncnn.load_state_dict(jt.load(path))\n\nfeat1_local, feat1_global = cnn(jittor_img1)\nfeat2_local, feat2_global = cnn(jittor_img2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Normalize the features\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def l2norm(node_feat):\n    return local_response_norm(\n        node_feat, node_feat.shape[1] * 2, alpha=node_feat.shape[1] * 2, beta=0.5, k=0)\n\nfeat1_local = l2norm(feat1_local)\nfeat1_global = l2norm(feat1_global)\nfeat2_local = l2norm(feat2_local)\nfeat2_global = l2norm(feat2_global)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Up-sample the features to the original image size and concatenate\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "feat1_local_upsample = jt.nn.interpolate(feat1_local, obj_resize, mode='bilinear')\nfeat1_global_upsample = jt.nn.interpolate(feat1_global, obj_resize, mode='bilinear')\nfeat2_local_upsample = jt.nn.interpolate(feat2_local, obj_resize, mode='bilinear')\nfeat2_global_upsample = jt.nn.interpolate(feat2_global, obj_resize, mode='bilinear')\nfeat1_upsample = jt.concat((feat1_local_upsample, feat1_global_upsample), dim=1)\nfeat2_upsample = jt.concat((feat2_local_upsample, feat2_global_upsample), dim=1)\nnum_features = feat1_upsample.shape[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extract node features by nearest interpolation\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rounded_kpts1 = jt.round(kpts1).long()\nrounded_kpts2 = jt.round(kpts2).long()\nnode1 = feat1_upsample[0, :, rounded_kpts1[1], rounded_kpts1[0]].t() # shape: NxC\nnode2 = feat2_upsample[0, :, rounded_kpts2[1], rounded_kpts2[0]].t() # shape: NxC"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Build edge features as edge lengths\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "kpts1_dis = (kpts1.unsqueeze(1) - kpts1.unsqueeze(2))\nkpts1_dis = jt.norm(kpts1_dis, p=2, dim=0).detach()\nkpts2_dis = (kpts2.unsqueeze(1) - kpts2.unsqueeze(2))\nkpts2_dis = jt.norm(kpts2_dis, p=2, dim=0).detach()\n\nQ1 = jt.exp(-kpts1_dis / obj_resize[0]).unsqueeze(-1).float32()\nQ2 = jt.exp(-kpts2_dis / obj_resize[0]).unsqueeze(-1).float32()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Call CIE matching model\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = pygm.cie(node1, node2, A1, A2, Q1, Q2, pretrain='voc')\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by CIE')\nax1 = plt.subplot(1, 2, 1)\nplot_image_with_graph(img1, kpts1, A1)\nax2 = plt.subplot(1, 2, 2)\nplot_image_with_graph(img2, kpts2, A2)\nidx, _ = jt.argmax(X, dim=1)\nfor i in range(X.shape[0]):\n    j = idx[i].item()\n    con = ConnectionPatch(xyA=kpts1[:, i], xyB=kpts2[:, j], coordsA=\"data\", coordsB=\"data\",\n                          axesA=ax1, axesB=ax2, color=\"red\" if i != j else \"green\")\n    plt.gca().add_artist(con)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training a deep graph matching model\nIn this section, we show how to build a deep graph matching model which supports end-to-end training.\nFor the image matching problem considered here, the model is composed of a CNN feature extractor and\na learnable matching module. Take the PCA-GM model as an example.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>This simple example is intended to show you how to do the basic execute and backward pass when\n    training an end-to-end deep graph matching neural network. A 'more formal' deep learning pipeline\n    should involve asynchronized data loader, batched operations, CUDA support and so on, which are\n    all omitted in consideration of simplicity. You may refer to [ThinkMatch](https://github.com/Thinklab-SJTU/ThinkMatch)\n    which is a research protocol with all these advanced features.</p></div>\n\nLet's firstly define the neural network model. By calling :func:`~pygmtools.utils.get_network`,\nit will simply return the network object.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class GMNet(jt.nn.Module):\n    def __init__(self):\n        super(GMNet, self).__init__()\n        self.gm_net = pygm.utils.get_network(pygm.pca_gm, pretrain=False) # fetch the network object\n        self.cnn = CNNNet(vgg16_cnn)\n\n    def execute(self, img1, img2, kpts1, kpts2, A1, A2):\n        # CNN feature extractor layers\n        feat1_local, feat1_global = self.cnn(img1)\n        feat2_local, feat2_global = self.cnn(img2)\n        feat1_local = l2norm(feat1_local)\n        feat1_global = l2norm(feat1_global)\n        feat2_local = l2norm(feat2_local)\n        feat2_global = l2norm(feat2_global)\n\n        # upsample feature map\n        feat1_local_upsample = jt.nn.interpolate(feat1_local, obj_resize, mode='bilinear')\n        feat1_global_upsample = jt.nn.interpolate(feat1_global, obj_resize, mode='bilinear')\n        feat2_local_upsample = jt.nn.interpolate(feat2_local, obj_resize, mode='bilinear')\n        feat2_global_upsample = jt.nn.interpolate(feat2_global, obj_resize, mode='bilinear')\n        feat1_upsample = jt.concat((feat1_local_upsample, feat1_global_upsample), dim=1)\n        feat2_upsample = jt.concat((feat2_local_upsample, feat2_global_upsample), dim=1)\n\n        # assign node features\n        rounded_kpts1 = jt.round(kpts1).long()\n        rounded_kpts2 = jt.round(kpts2).long()\n        node1 = feat1_upsample[0, :, rounded_kpts1[1], rounded_kpts1[0]].t()  # shape: NxC\n        node2 = feat2_upsample[0, :, rounded_kpts2[1], rounded_kpts2[0]].t()  # shape: NxC\n\n        # PCA-GM matching layers\n        X = pygm.pca_gm(node1, node2, A1, A2, network=self.gm_net) # the network object is reused\n        return X\n\nmodel = GMNet()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Define optimizer\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "optim = jt.optim.Adam(model.parameters(), lr=1e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Forward pass\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = model(jittor_img1, jittor_img2, kpts1, kpts2, A1, A2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Compute loss\nIn this example, the ground truth matching matrix is a diagonal matrix. We calculate the loss function via\n:func:`~pygmtools.utils.permutation_loss`\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X_gt = jt.init.eye(X.shape[0])\nloss = pygm.utils.permutation_loss(X, X_gt)\nprint(f'loss={loss:.4f}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Backward Pass\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "optim.backward(loss)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualize the gradients\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(4, 4))\nplt.title('Gradient Sizes of PCA-GM and VGG16 layers')\nplt.gca().set_xlabel('Layer Index')\nplt.gca().set_ylabel('Average Gradient Size')\ngrad_size = []\nfor param in model.parameters():\n    grad_size.append(jt.abs(param.opt_grad(optim)).mean().item())\nprint(grad_size)\nplt.stem(grad_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Update the model parameters. A deep learning pipeline should iterate the forward pass\nand backward pass steps until convergence.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "optim.step()\noptim.zero_grad()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>This example supports both GPU and CPU, and the online documentation is built by a CPU-only machine.\n    The efficiency will be significantly improved if you run this code on GPU.</p></div>\n\n\n"
      ]
    }
  ],
  "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.8.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}