tools/eval_corr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install the proper version of pytorch3d\n",
    "\n",
    "import sys\n",
    "import torch\n",
    "\n",
    "pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
    "version_str=\"\".join([\n",
    "    f\"py3{sys.version_info.minor}_cu\",\n",
    "    torch.version.cuda.replace(\".\",\"\"),\n",
    "    f\"_pyt{pyt_version_str}\"\n",
    "])\n",
    "!pip install fvcore\n",
    "!pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "\n",
    "import imageio\n",
    "import numpy as np\n",
    "import torch\n",
    "import tqdm\n",
    "from hydra import compose, initialize\n",
    "from hydra.utils import instantiate\n",
    "from omegaconf import OmegaConf\n",
    "from pytorch3d.renderer import MeshRasterizer, RasterizationSettings\n",
    "from pytorch3d.structures import Meshes\n",
    "from pytorch3d.utils import cameras_from_opencv_projection\n",
    "\n",
    "from tava.datasets.animal_parser import SubjectParser\n",
    "\n",
    "device = \"cuda:0\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# args for data\n",
    "ARGS_ANIMAL_WOLF = [\n",
    "    \"dataset=animal_wolf\", \"dataset.root_fp=/home/ruilongli/data/forest_and_friends_rendering/\",\n",
    "    \"hydra.run.dir=/home/ruilongli/workspace/TAVA/outputs/release/animal_wolf/Wolf_cub_full_RM_2/narf/\",\n",
    "]\n",
    "ARGS_ANIMAL_HARE=[\n",
    "    \"dataset=animal_hare\", \"dataset.root_fp=/home/ruilongli/data/forest_and_friends_rendering/\",\n",
    "    \"hydra.run.dir=/home/ruilongli/workspace/TAVA/outputs/release/animal_hare/Hare_male_full_RM/narf/\",\n",
    "]\n",
    "\n",
    "# args for method\n",
    "ARGS_TAVA_ANIMAL=[\"pos_enc=snarf\", \"loss_bone_w_mult=1.0\", \"pos_enc.offset_net_enabled=false\", \"model.shading_mode=null\"]\n",
    "ARGS_NARF=[\"pos_enc=narf\", \"model.shading_mode=null\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we set the arguments for ZJU_313 as an example.\n",
    "overrides = [\"resume=True\"] + ARGS_ANIMAL_HARE + ARGS_NARF\n",
    "split = \"val_ood\"\n",
    "\n",
    "# create the cfg\n",
    "with initialize(config_path=\"../configs\"):\n",
    "    cfg = compose(config_name=\"mipnerf_dyn\", overrides=overrides, return_hydra_config=True)\n",
    "    OmegaConf.resolve(cfg.hydra)\n",
    "    save_dir = cfg.hydra.run.dir\n",
    "    eval_imgs_dir = os.path.join(save_dir, \"eval_imgs\")\n",
    "\n",
    "# initialize dataset\n",
    "dataset = instantiate(\n",
    "    cfg.dataset, split=split, num_rays=None, cache_n_repeat=None,\n",
    ")\n",
    "meta_data_dict = {\n",
    "    action: dataset.parser.load_meta_data(action)\n",
    "    for action in dataset.parser.actions\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the ground-truth correspondances: pixel -> 3D canonical point\n",
    "\n",
    "image_size = int(800 * dataset.resize_factor)\n",
    "rasterizer = MeshRasterizer(\n",
    "    raster_settings=RasterizationSettings(image_size=image_size)\n",
    ")\n",
    "\n",
    "for index in tqdm.tqdm(dataset.index_list):\n",
    "    action, frame_id, camera_id = index\n",
    "    meta_id = dataset.encode_meta_id(action, frame_id)\n",
    "\n",
    "    K, c2w = dataset.parser.load_camera(action, frame_id, camera_id)\n",
    "    K = torch.from_numpy(K).float().to(device)\n",
    "    c2w = torch.from_numpy(c2w).float().to(device)\n",
    "\n",
    "    cameras = cameras_from_opencv_projection(\n",
    "        R=c2w.inverse()[None, :3, :3], \n",
    "        tvec=c2w.inverse()[None, :3, 3], \n",
    "        camera_matrix=K[None], \n",
    "        image_size=torch.tensor([[image_size, image_size]])\n",
    "    ).to(device)\n",
    "\n",
    "    faces = torch.from_numpy(\n",
    "        meta_data_dict[action][\"faces\"]).long().to(device)\n",
    "    verts = torch.from_numpy(\n",
    "        meta_data_dict[action][\"pose_verts\"][frame_id]).float().to(device)\n",
    "    rest_verts = torch.from_numpy(\n",
    "        meta_data_dict[action][\"rest_verts\"]).float().to(device)\n",
    "    meshes = Meshes(verts=[verts], faces=[faces])\n",
    "\n",
    "    fragments = rasterizer(meshes, cameras=cameras)\n",
    "\n",
    "    pix_to_face = fragments.pix_to_face\n",
    "    barycentric = fragments.bary_coords\n",
    "    pix_to_coord = torch.einsum(\n",
    "        \"bhwnvi,bhwnv->bhwni\", rest_verts[faces[pix_to_face]], barycentric)\n",
    "    masks = (pix_to_face != -1).squeeze(-1).squeeze(0)  # [h, w]\n",
    "    coords = pix_to_coord.squeeze(-2).squeeze(0) * masks[:, :, None]  # [h, w, 3]\n",
    "\n",
    "    image_to_save = torch.cat([coords, masks[:, :, None].float()], dim=-1)\n",
    "    image_path = os.path.join(\n",
    "        dataset.parser.root_dir, action, \"correspondence\", camera_id, \"%08d.exr\" % frame_id\n",
    "    )\n",
    "    os.makedirs(os.path.dirname(image_path), exist_ok=True)\n",
    "    imageio.imwrite(image_path, image_to_save.cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = torch.meshgrid(\n",
    "    torch.arange(800).long().to(device),  # X-Axis (columns)\n",
    "    torch.arange(800).long().to(device),  # Y-Axis (rows)\n",
    "    indexing=\"xy\",\n",
    ")\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def matching(warp_src, wrap_dst, chunk_size=8):\n",
    "    assert warp_src.dim() == wrap_dst.dim() == 2\n",
    "    errors, indices = [], []\n",
    "    for i in range(0, warp_src.shape[0], chunk_size):\n",
    "        warp_src_chunk = warp_src[i: i + chunk_size]\n",
    "        matching = torch.linalg.norm(\n",
    "            wrap_dst[:, None, :] - warp_src_chunk[None, :, :], dim=-1\n",
    "        ).min(dim=0)  # [chunk_size,]\n",
    "        errors.append(matching.values)\n",
    "        indices.append(matching.indices)\n",
    "    errors = torch.cat(errors, dim=0)\n",
    "    indices = torch.cat(indices, dim=0)\n",
    "    assert errors.shape[0] == warp_src.shape[0]\n",
    "    assert indices.shape[0] == indices.shape[0]\n",
    "    return errors, indices\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def matching_pairs(map_src, mask_src, map_dst, mask_dst, thre=1e-4):\n",
    "    warp_src = map_src[mask_src]\n",
    "    warp_dst = map_dst[mask_dst]\n",
    "    errors, indices = matching(warp_src, warp_dst)\n",
    "    selector = errors < thre\n",
    "    x_src = x[mask_src][selector]\n",
    "    y_src = y[mask_src][selector]\n",
    "    x_dst = x[mask_dst][indices][selector]\n",
    "    y_dst = y[mask_dst][indices][selector]\n",
    "    coord_src = torch.stack([x_src, y_src], dim=-1)\n",
    "    coord_dst = torch.stack([x_dst, y_dst], dim=-1)\n",
    "    return coord_src, coord_dst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the ground-truth correspondances: pixel -> pixel\n",
    "\n",
    "render_every = math.ceil(len(dataset) / 100)\n",
    "index_list = dataset.index_list[::render_every]\n",
    "print (\"index list\", len(index_list))\n",
    "\n",
    "pair_list = []\n",
    "for _, index_src in enumerate(index_list):\n",
    "    for _, index_dst in enumerate(index_list):\n",
    "        if index_src == index_dst: continue\n",
    "        pair_list.append([\n",
    "            dataset.index_list.index(index_src), index_src, \n",
    "            dataset.index_list.index(index_dst), index_dst\n",
    "        ])\n",
    "print (\"pair list\", len(pair_list))\n",
    "\n",
    "random.seed(42)\n",
    "random.shuffle(pair_list)\n",
    "for id1, index1, id2, index2 in tqdm.tqdm(pair_list[:2000]):\n",
    "    action1, frame_id1, camera_id1 = index1\n",
    "    action2, frame_id2, camera_id2 = index2\n",
    "\n",
    "    gt_map1 = torch.from_numpy(imageio.imread(os.path.join(\n",
    "        dataset.parser.root_dir, action1, \"correspondence\", camera_id1, \"%08d.exr\" % frame_id1\n",
    "    ))).to(device)\n",
    "    gt_map2 = torch.from_numpy(imageio.imread(os.path.join(\n",
    "        dataset.parser.root_dir, action2, \"correspondence\", camera_id2, \"%08d.exr\" % frame_id2\n",
    "    ))).to(device)\n",
    "\n",
    "    map_src = gt_map1[:, :, 0:3]\n",
    "    mask_src = gt_map1[:, :, 3] > 0.5\n",
    "    map_dst = gt_map2[:, :, 0:3]\n",
    "    mask_dst = gt_map2[:, :, 3] > 0.5\n",
    "    coord_src, coord_dst = matching_pairs(map_src, mask_src, map_dst, mask_dst, thre=1e-4)\n",
    "\n",
    "    cache_path = os.path.join(\n",
    "        \"/tmp\", \"tava_corr\", cfg.dataset.subject_id, split, \"gt\",\n",
    "        f\"{action1}_{frame_id1}_{camera_id1}___{action2}_{frame_id2}_{camera_id2}.npz\"\n",
    "    )\n",
    "    os.makedirs(os.path.dirname(cache_path), exist_ok=True)\n",
    "    np.savez(\n",
    "        cache_path, \n",
    "        coord_src = coord_src.cpu().numpy(), \n",
    "        coord_dst = coord_dst.cpu().numpy(), \n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = {\"p2p\": []}\n",
    "for id1, index1, id2, index2 in tqdm.tqdm(pair_list[:100]):\n",
    "    action1, frame_id1, camera_id1 = index1\n",
    "    action2, frame_id2, camera_id2 = index2\n",
    "\n",
    "    # get groundtruth\n",
    "    cache_path = os.path.join(\n",
    "        \"/tmp\", \"tava_corr\", cfg.dataset.subject_id, split, \"gt\",\n",
    "        f\"{action1}_{frame_id1}_{camera_id1}___{action2}_{frame_id2}_{camera_id2}.npz\"\n",
    "    )\n",
    "    gt_data = np.load(cache_path, allow_pickle=True)\n",
    "    gt_coord_src = torch.from_numpy(gt_data[\"coord_src\"]).to(device)\n",
    "    gt_coord_dst = torch.from_numpy(gt_data[\"coord_dst\"]).to(device)\n",
    "    if gt_coord_src.numel() == 0:\n",
    "        continue\n",
    "\n",
    "    # get predict map\n",
    "    meta_id1 = dataset.encode_meta_id(action1, frame_id1)\n",
    "    meta_id2 = dataset.encode_meta_id(action2, frame_id2)\n",
    "    # if args.method == \"ours\":\n",
    "    #     pred_map1 = torch.from_numpy(imageio.imread(\n",
    "    #         os.path.join(image_dir, f\"{id1:04d}_{sid1}_{fid1}_{cid1}.exr\"\n",
    "    #     ))).to(device)\n",
    "    #     pred_map2 = torch.from_numpy(imageio.imread(\n",
    "    #         os.path.join(image_dir, f\"{id2:04d}_{sid2}_{fid2}_{cid2}.exr\"\n",
    "    #     ))).to(device)\n",
    "    # else:\n",
    "    pred_map1 = torch.from_numpy(np.load(\n",
    "        os.path.join(\n",
    "            eval_imgs_dir, split, \n",
    "            f\"{id1:04d}_{cfg.dataset.subject_id}_{meta_id1}_{camera_id1}.npy\"\n",
    "        )\n",
    "    )).to(device)\n",
    "    pred_map2 = torch.from_numpy(np.load(\n",
    "        os.path.join(\n",
    "            eval_imgs_dir, split, \n",
    "            f\"{id2:04d}_{cfg.dataset.subject_id}_{meta_id2}_{camera_id2}.npy\"\n",
    "        )\n",
    "    )).to(device)\n",
    "\n",
    "    pred_map_src = pred_map1\n",
    "    pred_mask_src = torch.zeros_like(pred_map1[..., 0]).bool()\n",
    "    pred_mask_src[gt_coord_src[:, 1], gt_coord_src[:, 0]] = True\n",
    "    pred_map_dst = pred_map2\n",
    "    pred_mask_dst = pred_map2.sum(dim=-1) != 0\n",
    "    pred_coord_src, pred_coord_dst = matching_pairs(\n",
    "        pred_map_src, pred_mask_src, pred_map_dst, pred_mask_dst, thre=1e10\n",
    "    )\n",
    "    assert torch.allclose(pred_coord_src, gt_coord_src)\n",
    "    p2p = torch.linalg.norm(\n",
    "        pred_coord_dst.float() - gt_coord_dst.float(), dim=-1\n",
    "    ).mean()\n",
    "    metric[\"p2p\"].append(p2p)\n",
    "\n",
    "    cache_path = os.path.join(\n",
    "        \"/tmp\", \"tava_corr\", cfg.dataset.subject_id, split, \"pred\",\n",
    "        f\"{action1}_{frame_id1}_{camera_id1}___{action2}_{frame_id2}_{camera_id2}.npz\"\n",
    "    )\n",
    "    os.makedirs(os.path.dirname(cache_path), exist_ok=True)\n",
    "    np.savez(\n",
    "        cache_path, \n",
    "        {\n",
    "            \"coord_src\": pred_coord_src.cpu().numpy(), \n",
    "            \"coord_dst\": pred_coord_dst.cpu().numpy(), \n",
    "        }\n",
    "    )\n",
    "\n",
    "for key, value in metric.items():\n",
    "    metric[key] = sum(value) / len(value)\n",
    "\n",
    "print (metric)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f156ac58ad89dad164aee97f58386e335ecc1447b9ed93fd42b558e9d07d43f1"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('tava')",
   "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.7.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}