inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import random\n",
    "from argparse import Namespace\n",
    "from functools import partial\n",
    "from multiprocessing import Pool\n",
    "from pathlib import Path\n",
    "from typing import Union\n",
    "\n",
    "import numpy as np\n",
    "import kornia\n",
    "import torch\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms as T\n",
    "from PIL import Image\n",
    "from tqdm.auto import tqdm\n",
    "import pandas as pd\n",
    "from SoccerNet.Evaluation.utils_calibration import SoccerPitch\n",
    "\n",
    "\n",
    "from tvcalib.cam_modules import SNProjectiveCamera\n",
    "from tvcalib.module import TVCalibModule\n",
    "from tvcalib.cam_distr.tv_main_center import get_cam_distr, get_dist_distr\n",
    "from sn_segmentation.src.custom_extremities import generate_class_synthesis, get_line_extremities\n",
    "from tvcalib.sncalib_dataset import custom_list_collate, split_circle_central\n",
    "from tvcalib.utils.io import detach_dict, tensor2list\n",
    "from tvcalib.utils.objects_3d import SoccerPitchLineCircleSegments, SoccerPitchSNCircleCentralSplit\n",
    "from tvcalib.inference import InferenceDatasetCalibration, InferenceDatasetSegmentation, InferenceSegmentationModel\n",
    "from tvcalib.inference import get_camera_from_per_sample_output\n",
    "from tvcalib.utils import visualization_mpl_min as viz\n",
    "\n",
    "\n",
    "args = Namespace(\n",
    "        images_path=Path(\"data/datasets/wc14-test\"),\n",
    "        output_dir=Path(\"tmp\"),\n",
    "        checkpoint=\"data/segment_localization/train_59.pt\",\n",
    "        gpu=True,\n",
    "        nworkers=16,\n",
    "        batch_size_seg=16,\n",
    "        batch_size_calib=256,\n",
    "        image_width=1280,\n",
    "        image_height=720,\n",
    "        optim_steps=2000,\n",
    "        lens_dist=False,\n",
    "        write_masks=False\n",
    "    )\n",
    "device = \"cuda\" if args.gpu and torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "object3d = SoccerPitchLineCircleSegments(\n",
    "    device=device, base_field=SoccerPitchSNCircleCentralSplit()\n",
    ")\n",
    "object3dcpu = SoccerPitchLineCircleSegments(\n",
    "    device=\"cpu\", base_field=SoccerPitchSNCircleCentralSplit()\n",
    ")\n",
    "\n",
    "lines_palette = [0, 0, 0]\n",
    "for line_class in SoccerPitch.lines_classes:\n",
    "    lines_palette.extend(SoccerPitch.palette[line_class])\n",
    "\n",
    "fn_generate_class_synthesis = partial(generate_class_synthesis, radius=4)\n",
    "fn_get_line_extremities = partial(get_line_extremities, maxdist=30, width=455, height=256, num_points_lines=4, num_points_circles=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_seg = InferenceDatasetSegmentation(\n",
    "    args.images_path, args.image_width, args.image_height\n",
    ")\n",
    "print(\"number of images:\", len(dataset_seg))\n",
    "dataloader_seg = torch.utils.data.DataLoader(\n",
    "    dataset_seg,\n",
    "    batch_size=args.batch_size_seg,\n",
    "    num_workers=args.nworkers,\n",
    "    shuffle=False,\n",
    "    collate_fn=custom_list_collate,\n",
    ")\n",
    "\n",
    "model_seg = InferenceSegmentationModel(args.checkpoint, device)\n",
    "\n",
    "image_ids = []\n",
    "keypoints_raw = []\n",
    "(args.output_dir / \"masks\").mkdir(parents=True, exist_ok=True)\n",
    "for batch_dict in tqdm(dataloader_seg):\n",
    "    # semantic segmentation\n",
    "    # image_raw: [B, 3, image_height, image_width]\n",
    "    # image: [B, 3, 256, 455]\n",
    "    with torch.no_grad():\n",
    "        sem_lines = model_seg.inference(batch_dict[\"image\"].to(device))\n",
    "    sem_lines = sem_lines.cpu().numpy().astype(np.uint8)  # [B, 256, 455]\n",
    "\n",
    "    # point selection\n",
    "    with Pool(args.nworkers) as p:\n",
    "        skeletons_batch = p.map(fn_generate_class_synthesis, sem_lines)\n",
    "        keypoints_raw_batch = p.map(fn_get_line_extremities, skeletons_batch)\n",
    "\n",
    "    # write to file\n",
    "    if args.write_masks:\n",
    "        print(\"Write masks to file\")\n",
    "        for image_id, mask in zip(batch_dict[\"image_id\"], sem_lines):\n",
    "            mask = Image.fromarray(mask.astype(np.uint8)).convert(\"P\")\n",
    "            mask.putpalette(lines_palette)\n",
    "            mask.convert(\"RGB\").save(args.output_dir / \"masks\" / image_id)\n",
    "\n",
    "    image_ids.extend(batch_dict[\"image_id\"])\n",
    "    keypoints_raw.extend(keypoints_raw_batch)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_calib = TVCalibModule(\n",
    "    object3d,\n",
    "    get_cam_distr(1.96, args.batch_size_calib, 1),\n",
    "    get_dist_distr(args.batch_size_calib, 1) if args.lens_dist else None,\n",
    "    (args.image_height, args.image_width),\n",
    "    args.optim_steps,\n",
    "    device,\n",
    "    log_per_step=False,\n",
    "    tqdm_kwqargs=None,\n",
    ")\n",
    "\n",
    "dataset_calib = InferenceDatasetCalibration(keypoints_raw, args.image_width, args.image_height, object3d)\n",
    "dataloader_calib = torch.utils.data.DataLoader(dataset_calib, args.batch_size_calib, collate_fn=custom_list_collate)\n",
    "\n",
    "per_sample_output = defaultdict(list)\n",
    "per_sample_output[\"image_id\"] = [[x] for x in image_ids]\n",
    "for x_dict in dataloader_calib:\n",
    "    _batch_size = x_dict[\"lines__ndc_projected_selection_shuffled\"].shape[0]\n",
    "\n",
    "    points_line = x_dict[\"lines__px_projected_selection_shuffled\"]\n",
    "    points_circle = x_dict[\"circles__px_projected_selection_shuffled\"]\n",
    "    print(f\"{points_line.shape=}, {points_circle.shape=}\")\n",
    "\n",
    "    per_sample_loss, cam, _ = model_calib.self_optim_batch(x_dict)\n",
    "    output_dict = tensor2list(detach_dict({**cam.get_parameters(_batch_size), **per_sample_loss}))\n",
    "    \n",
    "    output_dict[\"points_line\"] = points_line\n",
    "    output_dict[\"points_circle\"] =  points_circle\n",
    "    for k in output_dict.keys():\n",
    "        per_sample_output[k].extend(output_dict[k])\n",
    "\n",
    "df = pd.DataFrame.from_dict(per_sample_output)\n",
    "\n",
    "df = df.explode(column=[k for k, v in per_sample_output.items() if isinstance(v, list)])\n",
    "df.set_index(\"image_id\", inplace=True, drop=False)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = df.iloc[0]\n",
    "\n",
    "image_id = Path(sample.image_id).stem\n",
    "print(f\"{image_id=}\")\n",
    "image = Image.open(args.images_path / sample.image_id).convert(\"RGB\")\n",
    "image = T.functional.to_tensor(image)\n",
    "\n",
    "cam = get_camera_from_per_sample_output(sample, args.lens_dist)\n",
    "print(cam, cam.str_lens_distortion_coeff(b=0) if args.lens_dist else \"\")\n",
    "points_line, points_circle = sample[\"points_line\"], sample[\"points_circle\"]\n",
    "\n",
    "if args.lens_dist:\n",
    "    # we visualize annotated points and image after undistortion\n",
    "    image = cam.undistort_images(image.unsqueeze(0).unsqueeze(0)).squeeze()\n",
    "    # print(points_line.shape) # expected: (1, 1, 3, S, N)\n",
    "    points_line = SNProjectiveCamera.static_undistort_points(points_line.unsqueeze(0).unsqueeze(0), cam).squeeze()\n",
    "    points_circle = SNProjectiveCamera.static_undistort_points(points_circle.unsqueeze(0).unsqueeze(0), cam).squeeze()\n",
    "else:\n",
    "    psi = None\n",
    "\n",
    "\n",
    "fig, ax = viz.init_figure(args.image_width, args.image_height)\n",
    "ax = viz.draw_image(ax, image)\n",
    "ax = viz.draw_reprojection(ax, object3dcpu, cam)\n",
    "ax = viz.draw_selected_points(\n",
    "    ax,\n",
    "    object3dcpu,\n",
    "    points_line,\n",
    "    points_circle,\n",
    "    kwargs_outer={\n",
    "        \"zorder\": 1000,\n",
    "        \"rasterized\": False,\n",
    "        \"s\": 500,\n",
    "        \"alpha\": 0.3,\n",
    "        \"facecolor\": \"none\",\n",
    "        \"linewidths\": 3,\n",
    "    },\n",
    "    kwargs_inner={\n",
    "        \"zorder\": 1000,\n",
    "        \"rasterized\": False,\n",
    "        \"s\": 50,\n",
    "        \"marker\": \".\",\n",
    "        \"color\": \"k\",\n",
    "        \"linewidths\": 4.0,\n",
    "    },\n",
    ")\n",
    "dpi = 50\n",
    "plt.savefig(args.output_dir / f\"{image_id}.pdf\", dpi=dpi)\n",
    "plt.savefig(args.output_dir / f\"{image_id}.svg\", dpi=dpi)\n",
    "plt.savefig(args.output_dir / f\"{image_id}.png\", dpi=dpi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "38404c93bbe569cd2e15f7b49ec04cc4c391d0e53654788e8fa6a229a378dae8"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}