test_time_denoising.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a442a7c2-abb4-48ad-9476-9bb0d0305e8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from scipy.signal import savgol_filter\n",
    "from tqdm.auto import tqdm\n",
    "import trimesh\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.data import Data\n",
    "from torch_cluster import nearest, radius_graph\n",
    "\n",
    "from diffusion_net.diffusion_net import DiffusionNet\n",
    "from diffusion_net.diffusion_operators import compute_diffusion_operators\n",
    "from losses import NCESoftmaxLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2e6a093f-b01e-4f38-8ffa-f594fd230aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "### UTILS ###\n",
    "\n",
    "def to_torch(array, device=\"cpu\"):\n",
    "    return torch.from_numpy(array).to(device)\n",
    "\n",
    "def read_mesh(file, process=False, is_np=False):\n",
    "    file = Path(file)\n",
    "    if file.suffix in [\".off\", \".ply\", \".obj\"]:\n",
    "        mesh = trimesh.load(file, process=process)\n",
    "        verts = np.array(mesh.vertices)\n",
    "        try:\n",
    "            faces = np.array(mesh.faces)\n",
    "        except AttributeError:\n",
    "            faces = np.zeros((0, 3))\n",
    "    else:\n",
    "        raise \"File extention not implemented yet!\"\n",
    "\n",
    "    if is_np:\n",
    "        return verts, faces\n",
    "\n",
    "    # return verts, faces\n",
    "    return torch.from_numpy(verts).float(), torch.from_numpy(faces)\n",
    "\n",
    "\n",
    "def pre_process(v, f):\n",
    "    v = torch.from_numpy(scale_unit_area(v.numpy(), f.numpy()))\n",
    "    v -= torch.mean(v, axis=0)\n",
    "    return v\n",
    "\n",
    "\n",
    "def mesh_area(verts, faces):\n",
    "    area = np.cross(verts[faces[:, 1]] - verts[faces[:, 0]], verts[faces[:, 2]] - verts[faces[:, 0]])\n",
    "    area = np.linalg.norm(area, axis=1) / 2\n",
    "    return area.sum()\n",
    "\n",
    "\n",
    "def scale_unit_area(verts, faces):\n",
    "    verts = np.copy(verts)\n",
    "    scale = 1 / np.sqrt(mesh_area(verts, faces))\n",
    "    return verts * scale"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1afc0c2f-c3fb-4ffa-9a8c-389d7ecfc8ed",
   "metadata": {},
   "source": [
    "## Test-Time Denoising (TTD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "178ce974-bb5f-486d-a078-343b882f1339",
   "metadata": {},
   "outputs": [],
   "source": [
    "neig = 100\n",
    "device = \"cuda:0\"\n",
    "lr = 1e-3\n",
    "total_iter = 400\n",
    "\n",
    "ROOT_DATA = Path(\"./data/animals\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dfd3bc08-0c46-4224-951e-525ec0ffaa64",
   "metadata": {},
   "outputs": [],
   "source": [
    "name1, name2 = \"68\", \"58\"\n",
    "\n",
    "## data loading\n",
    "# load shapes\n",
    "verts1, face1 = read_mesh(ROOT_DATA / f\"{name1}.off\")\n",
    "verts2, face2 = read_mesh(ROOT_DATA / f\"{name2}.off\")\n",
    "\n",
    "# gt maps\n",
    "map1, map2 = np.arange(3889), np.arange(3889)\n",
    "\n",
    "# preprocess\n",
    "verts1, verts2 = pre_process(verts1, face1), pre_process(verts2, face2)\n",
    "\n",
    "# load spectrals\n",
    "evals1, evecs1, evecs_trans1, A1, L1 = [to_torch(x) if isinstance(x, np.ndarray) else x for x in torch.load(ROOT_DATA / f\"{name1}.pt\")]\n",
    "edge_index1 = torch.LongTensor(np.vstack((L1.row, L1.col)))\n",
    "frames1, normals1, mass1, evals1, evecs1, grad_from_spectral1 = compute_diffusion_operators(verts1, face1, 128,  evals1, evecs1, A1, edge_index1)\n",
    "\n",
    "evals2, evecs2, evecs_trans2, A2, L2 = [to_torch(x) if isinstance(x, np.ndarray) else x for x in torch.load(ROOT_DATA / f\"{name2}.pt\")]\n",
    "edge_index2 = torch.LongTensor(np.vstack((L2.row, L2.col)))\n",
    "frames2, normals2, mass2, evals2, evecs2, grad_from_spectral2 = compute_diffusion_operators(verts2, face2, 128,  evals2, evecs2, A2, edge_index2)\n",
    "\n",
    "# load geodesic matrix\n",
    "geod_mat1 = np.load(ROOT_DATA / f\"{name1}.npy\").astype(float)\n",
    "geod_mat2 = np.load(ROOT_DATA / f\"{name2}.npy\").astype(float)\n",
    "\n",
    "# load noisy map\n",
    "T21_noisy = torch.from_numpy(np.loadtxt(ROOT_DATA / f\"{name2}_{name1}.map\").astype(int))\n",
    "T12_noisy = torch.from_numpy(np.loadtxt(ROOT_DATA / f\"{name1}_{name2}.map\").astype(int))\n",
    "\n",
    "# data object\n",
    "in_feat1, in_feat2 = verts1[None, ...].to(device), verts2[None, ...].to(device)\n",
    "N_IN = 3\n",
    "\n",
    "data1 = Data(pos=verts1[None, ...].to(device), verts=verts1[None, ...].to(device), x=in_feat1, face=face1.t()[None, ...].to(device), diffusion_evals=evals1[None, ...].to(device),\n",
    "             diffusion_evecs=evecs1[None, ...].to(device), diffusion_mass=mass1[None, ...].to(device), grad_from_spectral=grad_from_spectral1[None, ...].to(device),\n",
    "             evecs_trans=evecs_trans1[None, ...].to(device))\n",
    "\n",
    "data2 = Data(pos=verts2[None, ...].to(device), verts=verts2[None, ...].to(device), x=in_feat2, face=face2.t()[None, ...].to(device), diffusion_evals=evals2[None, ...].to(device),\n",
    "             diffusion_evecs=evecs2[None, ...].to(device), diffusion_mass=mass2[None, ...].to(device), grad_from_spectral=grad_from_spectral2[None, ...].to(device),\n",
    "             evecs_trans=evecs_trans2[None, ...].to(device))\n",
    "\n",
    "used_map21 = T21_noisy[None, ...].to(device)\n",
    "used_map12 = T12_noisy[None, ...].to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "9bbfad7c-927b-43b3-811d-ec50f48d9007",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model and optimizer\n",
    "model = DiffusionNet(dim_in=N_IN, dim_out=128).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.99))\n",
    "criterion = NCESoftmaxLoss(nce_t=0.07, nce_num_pairs=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "b01fb650-300c-4f14-b440-e26be2b47719",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2bf9d76d1c04a89bbc2825dedba1428",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/400 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# do test time optimization\n",
    "for i in tqdm(range(total_iter)):\n",
    "    shape1 = model(data1.clone())\n",
    "    shape2 = model(data2.clone())\n",
    "\n",
    "    shape1.features, shape2.features = F.normalize(shape1.features, p=2, dim=-1), F.normalize(shape2.features, p=2, dim=-1)\n",
    "    \n",
    "    loss = criterion(shape1.features, shape2.features, used_map21)\n",
    "    loss += criterion(shape2.features, shape1.features, used_map12)\n",
    "\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "T21_optim_nn = nearest(shape2.features[0], shape1.features[0]).detach().cpu()\n",
    "T12_optim_nn = nearest(shape1.features[0], shape2.features[0]).detach().cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "ecec3508-1292-4622-ac78-d76b8023413f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Map 2 -> 1\n",
      "Err noisy map: 5.51 | Err after TTD: 4.59\n"
     ]
    }
   ],
   "source": [
    "# print error\n",
    "err21_noisy = round(float(geod_mat1[map1, T21_noisy].mean()), 2)\n",
    "err21_ttd = round(float(geod_mat1[map1, T21_optim_nn].mean()), 2)\n",
    "print(\"Map 2 -> 1\")\n",
    "print(f\"Err noisy map: {err21_noisy} | Err after TTD: {err21_ttd}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "f413aef1-e7f8-47cc-af19-b8938af3f2da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Map 1 -> 2\n",
      "Err noisy map: 6.19 | Err after TTD: 4.55\n"
     ]
    }
   ],
   "source": [
    "# print error\n",
    "err12_noisy = round(float(geod_mat2[map2, T12_noisy].mean()), 2)\n",
    "err12_ttd = round(float(geod_mat2[map2, T12_optim_nn].mean()), 2)\n",
    "print(\"Map 1 -> 2\")\n",
    "print(f\"Err noisy map: {err12_noisy} | Err after TTD: {err12_ttd}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26fe5a98-5bdd-473f-86b9-d25c77c47d11",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geom_v1.13",
   "language": "python",
   "name": "geom_v1.13"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}