{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Matching Image Keypoints by QAP Solvers\n\nThis example shows how to match image keypoints by graph matching solvers provided by ``pygmtools``.\nThese solvers follow the Quadratic Assignment Problem formulation and can generally work out-of-box.\nThe matched images can be further processed for other downstream tasks.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Author: Runzhong Wang \n#\n# License: Mulan PSL v2 License" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

The following solvers support QAP formulation, and are included in this example:\n\n * :func:`~pygmtools.classic_solvers.rrwm` (classic solver)\n\n * :func:`~pygmtools.classic_solvers.ipfp` (classic solver)\n\n * :func:`~pygmtools.classic_solvers.sm` (classic solver)\n\n * :func:`~pygmtools.neural_solvers.ngm` (neural network solver)

\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import torch # pytorch backend\nimport torchvision # CV models\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 = 'pytorch' # set default backend for pygmtools" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 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 = torch.tensor(sio.loadmat('../data/willow_duck_0001.mat')['pts_coord'])\nkpts2 = torch.tensor(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)" ] }, { "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 torch.nonzero(A, as_tuple=False):\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 = torch.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": [ "We encode the length of edges as edge features\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "A1 = ((kpts1.unsqueeze(1) - kpts1.unsqueeze(2)) ** 2).sum(dim=0) * A1\nA1 = (A1 / A1.max()).to(dtype=torch.float32)\nA2 = ((kpts2.unsqueeze(1) - kpts2.unsqueeze(2)) ** 2).sum(dim=0) * A2\nA2 = (A2 / A2.max()).to(dtype=torch.float32)" ] }, { "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\nLet's adopt the VGG16 CNN model to extract node features.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "vgg16_cnn = torchvision.models.vgg16_bn(True)\ntorch_img1 = torch.from_numpy(np.array(img1, dtype=np.float32) / 256).permute(2, 0, 1).unsqueeze(0) # shape: BxCxHxW\ntorch_img2 = torch.from_numpy(np.array(img2, dtype=np.float32) / 256).permute(2, 0, 1).unsqueeze(0) # shape: BxCxHxW\nwith torch.set_grad_enabled(False):\n feat1 = vgg16_cnn.features(torch_img1)\n feat2 = vgg16_cnn.features(torch_img2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Normalize the features\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "num_features = feat1.shape[1]\ndef l2norm(node_feat):\n return torch.nn.functional.local_response_norm(\n node_feat, node_feat.shape[1] * 2, alpha=node_feat.shape[1] * 2, beta=0.5, k=0)\n\nfeat1 = l2norm(feat1)\nfeat2 = l2norm(feat2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Up-sample the features to the original image size\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "feat1_upsample = torch.nn.functional.interpolate(feat1, obj_resize, mode='bilinear')\nfeat2_upsample = torch.nn.functional.interpolate(feat2, obj_resize, mode='bilinear')" ] }, { "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 = torch.round(kpts1).to(dtype=torch.long)\nrounded_kpts2 = torch.round(kpts2).to(dtype=torch.long)\nnode1 = feat1_upsample[0, :, rounded_kpts1[0], rounded_kpts1[1]].t() # shape: NxC\nnode2 = feat2_upsample[0, :, rounded_kpts2[0], rounded_kpts2[1]].t() # shape: NxC" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Build affinity matrix\nWe follow the formulation of Quadratic Assignment Problem (QAP):\n\n\\begin{align}&\\max_{\\mathbf{X}} \\ \\texttt{vec}(\\mathbf{X})^\\top \\mathbf{K} \\texttt{vec}(\\mathbf{X})\\\\\n s.t. \\quad &\\mathbf{X} \\in \\{0, 1\\}^{n_1\\times n_2}, \\ \\mathbf{X}\\mathbf{1} = \\mathbf{1}, \\ \\mathbf{X}^\\top\\mathbf{1} \\leq \\mathbf{1}\\end{align}\n\nwhere the first step is to build the affinity matrix ($\\mathbf{K}$)\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "conn1, edge1 = pygm.utils.dense_to_sparse(A1)\nconn2, edge2 = pygm.utils.dense_to_sparse(A2)\nimport functools\ngaussian_aff = functools.partial(pygm.utils.gaussian_aff_fn, sigma=1) # set affinity function\nK = pygm.utils.build_aff_mat(node1, edge1, conn1, node2, edge2, conn2, edge_aff_fn=gaussian_aff)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Visualization of the affinity matrix. For graph matching problem with $N$ nodes, the affinity matrix\nhas $N^2\\times N^2$ elements because there are $N^2$ edges in each graph.\n\n

Note

The diagonal elements are node affinities, the off-diagonal elements are edge features.

\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plt.figure(figsize=(4, 4))\nplt.title(f'Affinity Matrix (size: {K.shape[0]}$\\\\times${K.shape[1]})')\nplt.imshow(K.numpy(), cmap='Blues')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Solve graph matching problem by RRWM solver\nSee :func:`~pygmtools.classic_solvers.rrwm` for the API reference.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.rrwm(K, kpts1.shape[1], kpts2.shape[1])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The output of RRWM is a soft matching matrix. Hungarian algorithm is then adopted to reach a discrete matching matrix\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.hungarian(X)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plot the matching\nThe correct matchings are marked by green, and wrong matchings are marked by red. In this example, the nodes are\nordered by their ground truth classes (i.e. the ground truth matching matrix is a diagonal matrix).\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by RRWM')\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)\nfor i in range(X.shape[0]):\n j = torch.argmax(X[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": [ "## Solve by other solvers\nWe could also do a quick benchmarking of other solvers on this specific problem.\n\n### IPFP solver\nSee :func:`~pygmtools.classic_solvers.ipfp` for the API reference.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.ipfp(K, kpts1.shape[1], kpts2.shape[1])\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by IPFP')\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)\nfor i in range(X.shape[0]):\n j = torch.argmax(X[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": [ "### SM solver\nSee :func:`~pygmtools.classic_solvers.sm` for the API reference.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.sm(K, kpts1.shape[1], kpts2.shape[1])\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by SM')\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)\nfor i in range(X.shape[0]):\n j = torch.argmax(X[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": [ "### NGM solver\nSee :func:`~pygmtools.neural_solvers.ngm` for the API reference.\n\n

Note

The NGM solvers are pretrained on a different problem setting, so their performance may seem inferior.\n To improve their performance, you may change the way of building affinity matrices, or try finetuning\n NGM on the new problem.

\n\nThe NGM solver pretrained on Willow dataset:\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.ngm(K, kpts1.shape[1], kpts2.shape[1], pretrain='willow')\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by NGM (willow pretrain)')\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)\nfor i in range(X.shape[0]):\n j = torch.argmax(X[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": [ "The NGM solver pretrained on VOC dataset:\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "X = pygm.ngm(K, kpts1.shape[1], kpts2.shape[1], pretrain='voc')\nX = pygm.hungarian(X)\n\nplt.figure(figsize=(8, 4))\nplt.suptitle('Image Matching Result by NGM (voc pretrain)')\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)\nfor i in range(X.shape[0]):\n j = torch.argmax(X[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)" ] } ], "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.9.7" } }, "nbformat": 4, "nbformat_minor": 0 }