test_padding.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random integer image:\n",
      "[[[[ -6.  -6.  -1.  -2.   4.]\n",
      "   [  3.  -2. -11.   2.  -1.]\n",
      "   [  7.   1.   1.   4.  -1.]\n",
      "   [ -1.  -3.   6.  10.   0.]\n",
      "   [  3.  -2.  -4. -12.  -1.]]]]\n",
      "\n",
      "Random integer convolution kernel:\n",
      "[[[[-4.  3.  1. -1. -3.]\n",
      "   [ 5.  2. -3.  0.  6.]\n",
      "   [ 6. -7. 13. -2.  2.]\n",
      "   [ 8. 10. -2. -2. -5.]\n",
      "   [ 6. -1.  4.  2. -4.]]]]\n",
      "\n",
      "Lagrange polynomial padding:\n",
      "[[[[17820.  4010.   144.    97.   306.  -230.    50.  5269. 22112.]\n",
      "   [ 4738.  1049.    18.    13.    81.   -52.    19.  1378.  5788.]\n",
      "   [  173.    34.    -6.    -6.    -1.    -2.     4.    54.   209.]\n",
      "   [ -423.   -86.     3.    -2.   -11.     2.    -1.  -122.  -527.]\n",
      "   [    9.    14.     7.     1.     1.     4.    -1.   -33.  -119.]\n",
      "   [  135.    35.    -1.    -3.     6.    10.     0.   -26.   -63.]\n",
      "   [  228.    54.     3.    -2.    -4.   -12.    -1.    88.   348.]\n",
      "   [ 2168.   524.    74.    34.   -16.  -132.    -6.  1034.  4024.]\n",
      "   [ 9442.  2274.   298.   153.     4.  -458.   -16.  4073. 16078.]]]]\n",
      "Lagrange polynomial padded, then convolved:\n",
      "[[[[-37294. -10606.    177.  -6228. -34906.]\n",
      "   [-16081.  -4098.     13.  -3570. -17936.]\n",
      "   [  -156.     62.    -15.  -1194.  -4813.]\n",
      "   [ 17394.   4883.    220.  -4767. -15136.]\n",
      "   [ 80126.  21810.   2005. -22072. -78736.]]]]\n",
      "\n",
      "Diff convolved:\n",
      "[[[[-37294. -10606.    177.  -6228. -34906.]\n",
      "   [-16081.  -4098.     13.  -3570. -17936.]\n",
      "   [  -156.     62.    -15.  -1194.  -4813.]\n",
      "   [ 17394.   4883.    220.  -4767. -15136.]\n",
      "   [ 80126.  21810.   2005. -22072. -78736.]]]]\n",
      "\n",
      "Difference between Lagrange polynomial padded convolved and diff convolved:\n",
      "[[[[0. 0. 0. 0. 0.]\n",
      "   [0. 0. 0. 0. 0.]\n",
      "   [0. 0. 0. 0. 0.]\n",
      "   [0. 0. 0. 0. 0.]\n",
      "   [0. 0. 0. 0. 0.]]]]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from experiments.exp_filtering_methods import conv2d_diff\n",
    "\n",
    "lagrange_padding_coefs = [torch.tensor(c, dtype=torch.float32) for c in [\n",
    "    [[1.0, -3.0, 3.0]],\n",
    "    [[1.0, -5.0, 10.0, -10.0, 5.0],\n",
    "     [5.0, -24.0, 45.0, -40.0, 15.0]],\n",
    "    [[1.0, -7.0, 21.0, -35.0, 35.0, -21.0, 7.0],\n",
    "     [7.0, -48.0, 140.0, -224.0, 210.0, -112.0, 28.0],\n",
    "     [28.0, -189.0, 540.0, -840.0, 756.0, -378.0, 84.0]]\n",
    "]]\n",
    "\n",
    "def conv2d_lagrange_extrap(image, kernel, returns_padded=False):\n",
    "    \"\"\" extrapolation-based conv2d methods \"\"\"\n",
    "    # pad with zero\n",
    "    k_size = kernel.shape[2]\n",
    "    p = k_size // 2\n",
    "    padded = F.pad(image, pad=(p, p, p, p))\n",
    "\n",
    "    # extrapolation\n",
    "    # top\n",
    "    padded[:, :, :p, p:-p] = torch.einsum(\n",
    "        '...yx,sy->sx', image[:, :, :k_size, :], lagrange_padding_coefs[k_size//2 - 1].flip(dims=(0, 1))\n",
    "    )\n",
    "    # bottom\n",
    "    padded[:, :, -p:, p:-p] = torch.einsum(\n",
    "        '...yx,sy->sx', image[:, :, -k_size:, :], lagrange_padding_coefs[k_size//2 - 1]\n",
    "    )\n",
    "    # left\n",
    "    padded[:, :, :, :p] = torch.einsum(\n",
    "        '...yx,sx->ys', padded[:, :, :, p:p+k_size], lagrange_padding_coefs[k_size//2 - 1].flip(dims=(0, 1))\n",
    "    )\n",
    "    # right\n",
    "    padded[:, :, :, -p:] = torch.einsum(\n",
    "        '...yx,sx->ys', padded[:, :, :, -p-k_size:-p], lagrange_padding_coefs[k_size//2 - 1]\n",
    "    )\n",
    "    \n",
    "    if returns_padded:\n",
    "        return F.conv2d(padded, kernel), padded\n",
    "    return F.conv2d(padded, kernel)\n",
    "\n",
    "\n",
    "M, N, = 5, 5\n",
    "\n",
    "torch.manual_seed(0)\n",
    "image = torch.round(torch.randn((1, 1, M, M))*5)\n",
    "kernel = torch.round(torch.randn((1, 1, N, N))*5)\n",
    "\n",
    "np.set_printoptions(suppress=True)\n",
    "torch.set_printoptions(sci_mode=False)\n",
    "\n",
    "print(\"Random integer image:\")\n",
    "print(image.numpy())\n",
    "\n",
    "print(\"\\nRandom integer convolution kernel:\")\n",
    "print(kernel.numpy())\n",
    "\n",
    "print(\"\\nLagrange polynomial padding:\")\n",
    "lagrange_padded_convolved, lagrange_padded = conv2d_lagrange_extrap(image, kernel, returns_padded=True)\n",
    "print(lagrange_padded.numpy())\n",
    "print(\"Lagrange polynomial padded, then convolved:\")\n",
    "print(lagrange_padded_convolved.numpy())\n",
    "\n",
    "diff_convolved = conv2d_diff(image, kernel)\n",
    "print(\"\\nDiff convolved:\")\n",
    "print(diff_convolved.numpy())\n",
    "\n",
    "print(\"\\nDifference between Lagrange polynomial padded convolved and diff convolved:\")\n",
    "print((diff_convolved - lagrange_padded_convolved).numpy())\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}