knn_matting_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xSPdfCjyLa7z",
    "outputId": "888ba494-84db-4a83-a849-cde6d611f62e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⏬ Downloading https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh...\n",
      "📦 Installing...\n",
      "📌 Adjusting configuration...\n",
      "🩹 Patching environment...\n",
      "⏲ Done in 0:00:29\n",
      "🔁 Restarting kernel...\n"
     ]
    }
   ],
   "source": [
    "!pip install -q condacolab\n",
    "import condacolab\n",
    "condacolab.install()\n",
    "!conda install -c pytorch faiss-gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LiQiX8-qKt_J",
    "outputId": "3db80b26-dce7-4288-bcd7-f3f0e0f28890"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting package metadata (current_repodata.json): - \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\bdone\n",
      "Solving environment: | \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\b- \b\b\\ \b\b| \b\b/ \b\bdone\n",
      "\n",
      "## Package Plan ##\n",
      "\n",
      "  environment location: /usr/local\n",
      "\n",
      "  added / updated specs:\n",
      "    - faiss-gpu\n",
      "\n",
      "\n",
      "The following packages will be downloaded:\n",
      "\n",
      "    package                    |            build\n",
      "    ---------------------------|-----------------\n",
      "    ca-certificates-2021.10.8  |       ha878542_0         139 KB  conda-forge\n",
      "    certifi-2021.10.8          |   py37h89c1867_1         145 KB  conda-forge\n",
      "    conda-4.10.3               |   py37h89c1867_3         3.1 MB  conda-forge\n",
      "    cudatoolkit-11.1.1         |       h6406543_8        1.20 GB  conda-forge\n",
      "    faiss-1.7.1                |py37cuda111h7f21d35_0_cuda         2.0 MB  conda-forge\n",
      "    faiss-gpu-1.7.1            |       h788eb59_0          15 KB  conda-forge\n",
      "    libblas-3.9.0              |11_linux64_openblas          12 KB  conda-forge\n",
      "    libcblas-3.9.0             |11_linux64_openblas          11 KB  conda-forge\n",
      "    libfaiss-1.7.1             |cuda111hf54f04a_0_cuda        79.8 MB  conda-forge\n",
      "    libfaiss-avx2-1.7.1        |cuda111h1234567_0_cuda        79.9 MB  conda-forge\n",
      "    libgfortran-ng-11.2.0      |      h69a702a_11          19 KB  conda-forge\n",
      "    libgfortran5-11.2.0        |      h5c6108e_11         1.7 MB  conda-forge\n",
      "    liblapack-3.9.0            |11_linux64_openblas          11 KB  conda-forge\n",
      "    libopenblas-0.3.17         |pthreads_h8fe5266_1         9.2 MB  conda-forge\n",
      "    numpy-1.20.3               |   py37h038b26d_1         5.7 MB  conda-forge\n",
      "    python_abi-3.7             |          2_cp37m           4 KB  conda-forge\n",
      "    ------------------------------------------------------------\n",
      "                                           Total:        1.38 GB\n",
      "\n",
      "The following NEW packages will be INSTALLED:\n",
      "\n",
      "  cudatoolkit        conda-forge/linux-64::cudatoolkit-11.1.1-h6406543_8\n",
      "  faiss              conda-forge/linux-64::faiss-1.7.1-py37cuda111h7f21d35_0_cuda\n",
      "  faiss-gpu          conda-forge/linux-64::faiss-gpu-1.7.1-h788eb59_0\n",
      "  libblas            conda-forge/linux-64::libblas-3.9.0-11_linux64_openblas\n",
      "  libcblas           conda-forge/linux-64::libcblas-3.9.0-11_linux64_openblas\n",
      "  libfaiss           conda-forge/linux-64::libfaiss-1.7.1-cuda111hf54f04a_0_cuda\n",
      "  libfaiss-avx2      conda-forge/linux-64::libfaiss-avx2-1.7.1-cuda111h1234567_0_cuda\n",
      "  libgfortran-ng     conda-forge/linux-64::libgfortran-ng-11.2.0-h69a702a_11\n",
      "  libgfortran5       conda-forge/linux-64::libgfortran5-11.2.0-h5c6108e_11\n",
      "  liblapack          conda-forge/linux-64::liblapack-3.9.0-11_linux64_openblas\n",
      "  libopenblas        conda-forge/linux-64::libopenblas-0.3.17-pthreads_h8fe5266_1\n",
      "  numpy              conda-forge/linux-64::numpy-1.20.3-py37h038b26d_1\n",
      "\n",
      "The following packages will be UPDATED:\n",
      "\n",
      "  ca-certificates                      2020.12.5-ha878542_0 --> 2021.10.8-ha878542_0\n",
      "  certifi                          2020.12.5-py37h89c1867_1 --> 2021.10.8-py37h89c1867_1\n",
      "  conda                                4.9.2-py37h89c1867_0 --> 4.10.3-py37h89c1867_3\n",
      "  python_abi                                    3.7-1_cp37m --> 3.7-2_cp37m\n",
      "\n",
      "\n",
      "\n",
      "Downloading and Extracting Packages\n",
      "libfaiss-avx2-1.7.1  | 79.9 MB   | : 100% 1.0/1 [00:24<00:00, 24.01s/it]               \n",
      "certifi-2021.10.8    | 145 KB    | : 100% 1.0/1 [00:00<00:00, 20.41it/s]\n",
      "cudatoolkit-11.1.1   | 1.20 GB   | : 100% 1.0/1 [02:30<00:00, 150.43s/it]              \n",
      "python_abi-3.7       | 4 KB      | : 100% 1.0/1 [00:00<00:00, 24.43it/s]\n",
      "faiss-1.7.1          | 2.0 MB    | : 100% 1.0/1 [00:01<00:00,  1.90s/it]\n",
      "numpy-1.20.3         | 5.7 MB    | : 100% 1.0/1 [00:01<00:00,  1.11s/it]\n",
      "libgfortran5-11.2.0  | 1.7 MB    | : 100% 1.0/1 [00:00<00:00,  2.97it/s]\n",
      "libgfortran-ng-11.2. | 19 KB     | : 100% 1.0/1 [00:00<00:00, 28.87it/s]\n",
      "libblas-3.9.0        | 12 KB     | : 100% 1.0/1 [00:00<00:00, 30.52it/s]\n",
      "libcblas-3.9.0       | 11 KB     | : 100% 1.0/1 [00:00<00:00, 30.44it/s]\n",
      "liblapack-3.9.0      | 11 KB     | : 100% 1.0/1 [00:00<00:00, 34.92it/s]\n",
      "libfaiss-1.7.1       | 79.8 MB   | : 100% 1.0/1 [00:23<00:00, 23.49s/it]               \n",
      "faiss-gpu-1.7.1      | 15 KB     | : 100% 1.0/1 [00:00<00:00,  2.97it/s]\n",
      "conda-4.10.3         | 3.1 MB    | : 100% 1.0/1 [00:00<00:00,  1.79it/s]\n",
      "libopenblas-0.3.17   | 9.2 MB    | : 100% 1.0/1 [00:01<00:00,  1.46s/it]              \n",
      "ca-certificates-2021 | 139 KB    | : 100% 1.0/1 [00:00<00:00, 18.92it/s]\n",
      "Preparing transaction: \\ \b\bdone\n",
      "Verifying transaction: / \b\b- \b\bdone\n",
      "Executing transaction: | \b\b/ \b\b- \b\b\\ \b\b| By downloading and using the CUDA Toolkit conda packages, you accept the terms and conditions of the CUDA End User License Agreement (EULA): https://docs.nvidia.com/cuda/eula/index.html\n",
      "\n",
      "\b\bdone\n"
     ]
    }
   ],
   "source": [
    "!conda install -c pytorch faiss-gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WKKD1V1ikQfv",
    "outputId": "1090b0c9-de06-4b6a-f854-8da5e290a2fa"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'AlphaMatting-Information-Flow'...\n",
      "remote: Enumerating objects: 11, done.\u001b[K\n",
      "remote: Total 11 (delta 0), reused 0 (delta 0), pack-reused 11\u001b[K\n",
      "Unpacking objects: 100% (11/11), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/jinczing/info-flow-matting.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "2WSms_XHKZtT"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sklearn.neighbors as neighbors\n",
    "import scipy.sparse as sparse\n",
    "import warnings\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import time\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import faiss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "3X81quDTKfJr"
   },
   "outputs": [],
   "source": [
    "def knn_matting(image, trimap, my_lambda=100, use_faiss=False, use_weak_spatial=False,\n",
    "                use_vgg_feature=False, use_bicg=False):\n",
    "    timer = time.time()\n",
    "    [h, w, c] = image.shape\n",
    "    image, trimap = image / 255.0, trimap / 255.0\n",
    "    foreground = (trimap == 1.0).astype(int)\n",
    "    background = (trimap == 0.0).astype(int)\n",
    "    all_constraints = (foreground + background)\n",
    "\n",
    "    ####################################################\n",
    "    # TODO: find KNN for the given image\n",
    "    ####################################################\n",
    "    k = 11\n",
    "    k2 = 5\n",
    "    k3 = 5\n",
    "    lamb = 100\n",
    "    C = 3+2\n",
    "    f_lamb = 1\n",
    "    norm = np.sqrt(h*h+w*w)\n",
    "    u, v = np.meshgrid(np.arange(0, w)/norm, np.arange(0, h)/norm)\n",
    "    X = np.concatenate([image[..., :],\n",
    "               np.repeat(np.expand_dims(u+np.random.rand(h, w)*1e-6, axis=-1), 1, axis=-1), # +np.random.rand(h, w)*1e-6\n",
    "               np.repeat(np.expand_dims(v+np.random.rand(h, w)*1e-6, axis=-1), 1, axis=-1)], axis=-1) # +np.random.rand(h, w)*1e-6\n",
    "    if use_vgg_feature:\n",
    "      preprocess = transforms.Compose([\n",
    "          transforms.ToTensor(),\n",
    "          transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "      ])\n",
    "      image = preprocess(image).unsqueeze(0).float()\n",
    "      model = torchvision.models.vgg16(pretrained=True)\n",
    "      model.eval()\n",
    "      device = 'cuda'\n",
    "      model = model.to(device)\n",
    "      image = image.to(device)\n",
    "      f = model.features[:3](image)\n",
    "      f = f.detach().cpu().numpy().squeeze().transpose((1, 2, 0))\n",
    "      f = f.reshape(-1, 64)\n",
    "\n",
    "    X = X.reshape(-1, C).astype('float32')\n",
    "    print(X.shape)\n",
    "    if use_faiss:\n",
    "      res = faiss.StandardGpuResources()  # use a single GPU\n",
    "\n",
    "      # first knn\n",
    "      timer = time.time()\n",
    "      index = faiss.IndexFlatL2(C)\n",
    "      index = faiss.index_cpu_to_gpu(res, 0, index)\n",
    "      index.add(X)\n",
    "      _, knns = index.search(X, k)\n",
    "      print('first knn time', time.time()-timer)\n",
    "\n",
    "      # second knn (weak spatial constraints)\n",
    "      if use_weak_spatial:\n",
    "        timer = time.time()\n",
    "        index = faiss.IndexFlatL2(C)\n",
    "        index = faiss.index_cpu_to_gpu(res, 0, index)\n",
    "        X[:, -2:] = X[:, -2:]/100\n",
    "        index.add(X)\n",
    "        _, knns2 = index.search(X, k2)\n",
    "        print('second knn time', time.time()-timer)\n",
    "\n",
    "        if use_vgg_feature:\n",
    "          # third knn (vgg features)\n",
    "          timer = time.time()\n",
    "          index = faiss.IndexFlatL2(64)\n",
    "          index = faiss.index_cpu_to_gpu(res, 0, index)\n",
    "\n",
    "          index.add(np.ascontiguousarray(f))\n",
    "          _, knns3 = index.search(np.ascontiguousarray(f), k3)\n",
    "          knns3 = knns3.reshape((h*w, k3))\n",
    "          \n",
    "          print('third knn time', time.time()-timer)\n",
    "\n",
    "          knns = np.concatenate((knns, knns2, knns3), axis=-1).astype(int)\n",
    "        else:\n",
    "          knns = np.concatenate((knns, knns2), axis=-1).astype(int)\n",
    "          k3 = 0\n",
    "      else:\n",
    "        k2 = 0\n",
    "        k3 = 0\n",
    "    else:\n",
    "      # first knn\n",
    "      timer = time.time()\n",
    "      knn = neighbors.NearestNeighbors(n_neighbors=k, p=1).fit(X)\n",
    "      knns = knn.kneighbors(X)[1]\n",
    "      print('first knn time', time.time()-timer)\n",
    "\n",
    "      # second knn\n",
    "      if use_weak_spatial:\n",
    "        timer = time.time()\n",
    "        X[:, -2:] = X[:, -2:]/100\n",
    "        knn = neighbors.NearestNeighbors(n_neighbors=k2, p=1).fit(X)\n",
    "        knns2 = knn.kneighbors(X)[1]\n",
    "        print('second knn time', time.time()-timer)\n",
    "\n",
    "        if use_vgg_feature:\n",
    "          # third knn\n",
    "          timer = time.time()\n",
    "          knn = neighbors.NearestNeighbors(n_neighbors=k3, p=1).fit(f)\n",
    "          knns3 = knn.kneighbors(f)[1]\n",
    "          print('third knn time', time.time()-timer)\n",
    "\n",
    "          knns = np.concatenate((knns, knns2, knns3), axis=-1)\n",
    "        else:\n",
    "          knns = np.concatenate((knns, knns2), axis=-1)\n",
    "          k3 = 0\n",
    "      else:\n",
    "        k2 = 0\n",
    "        k3 = 0\n",
    "\n",
    "    print(knns.shape)\n",
    "\n",
    "    ####################################################\n",
    "    # TODO: compute the affinity matrix A\n",
    "    #       and all other stuff needed\n",
    "    ####################################################\n",
    "    row_inds = np.repeat(np.arange(h*w), k+k2+k3)\n",
    "    col_inds = knns.reshape(h*w*(k+k2+k3))\n",
    "    A = 1 - np.linalg.norm(X[row_inds, :C] - X[col_inds, :C], axis=1, ord=1)/C\n",
    "    A = sparse.coo_matrix((A, (row_inds, col_inds)), shape=(h*w, h*w))\n",
    "\n",
    "    D_script = sparse.diags(np.ravel(A.sum(axis=1)))\n",
    "    L = D_script-A\n",
    "    D = sparse.diags(np.ravel(all_constraints[:,:]))\n",
    "    v = np.ravel(foreground[:,:])\n",
    "    c = lamb*np.transpose(v)\n",
    "    H = (L + lamb*D)\n",
    "    \n",
    "    ####################################################\n",
    "    # TODO: solve for the linear system,\n",
    "    #       note that you may encounter en error\n",
    "    #       if no exact solution exists\n",
    "    ####################################################\n",
    "    print('start solving')\n",
    "    timer = time.time()\n",
    "    if use_bicg:\n",
    "      alpha, _ = sparse.linalg.bicg(H, c, atol=1e-10, maxiter=6000)\n",
    "      print('solving time:', time.time()-timer)\n",
    "      alpha = alpha.reshape(h, w)\n",
    "    else:\n",
    "      warnings.filterwarnings('error')\n",
    "      try:\n",
    "        alpha = sparse.linalg.spsolve(H, c)\n",
    "      except Warning:\n",
    "        print('using bicg to solve instead')\n",
    "        alpha, _ = sparse.linalg.bicg(H, c, atol=1e-10, maxiter=6000)\n",
    "      print('solving time:', time.time()-timer)\n",
    "      alpha = alpha.reshape(h, w)\n",
    "    print(alpha.min(), alpha.max())\n",
    "\n",
    "    return alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Nxkgu-X2Kk5x",
    "outputId": "7b063087-9a9b-4a1d-eccb-dc0da3eabda8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(143640, 5)\n",
      "first knn time 0.6080746650695801\n",
      "second knn time 0.4530041217803955\n",
      "(143640, 16)\n",
      "start solving\n",
      "solving time: 0.6991984844207764\n",
      "-3.116766201295849e-06 1.001473862033135\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# run knn matting\n",
    "image = cv2.imread('./image/bear.png')\n",
    "trimap = cv2.imread('./trimap/bear.png', cv2.IMREAD_GRAYSCALE)\n",
    "background = cv2.imread('./image/forest.jpg')\n",
    "background = cv2.resize(background, (image.shape[1], image.shape[0]))\n",
    "\n",
    "alpha = knn_matting(image, trimap, use_faiss=True, use_weak_spatial=True, use_vgg_feature=False, use_bicg=True)\n",
    "alpha = alpha[:, :, np.newaxis].clip(0, 1)\n",
    "cv2.imwrite('./result/bear_alpha_faiss.png', alpha*255)\n",
    "\n",
    "####################################################\n",
    "# TODO: pick up your own background image, \n",
    "#       and merge it with the foreground\n",
    "####################################################\n",
    "result = (1-alpha)*background + alpha*image\n",
    "cv2.imwrite('./result/bear_faiss.png', result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rJBlbzYmdyx1",
    "outputId": "2a1f5254-e3d1-4db1-d5b6-1b11bfc91490"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "134155\n",
      "(9485, 14, 1)\n",
      "(9485, 14, 3)\n",
      "(9485, 5, 5)\n",
      "1.3548182723431954 -0.20146884252984487\n",
      "info flow time 26.96239733695984\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# run information flow matting\n",
    "sys.path.append('./info_flow/')\n",
    "from info_flow.alpha import info_flow\n",
    "image = cv2.imread('./image/bear.png')\n",
    "trimap = cv2.imread('./trimap/bear.png', cv2.IMREAD_GRAYSCALE)\n",
    "background = cv2.imread('./image/forest.jpg')\n",
    "background = cv2.resize(background, (image.shape[1], image.shape[0]))\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
    "timer = time.time()\n",
    "alpha = info_flow(image, trimap)\n",
    "print('info flow time', time.time()-timer)\n",
    "alpha = alpha[:, :, np.newaxis].clip(0, 1)\n",
    "cv2.imwrite('./result/bear_alpha_info.png', alpha*255)\n",
    "\n",
    "result = (1-alpha)*background + alpha*image\n",
    "cv2.imwrite('./result/bear_info.png', result)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "knn-matting-gpu.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}