{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import h5py as h5\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as mpatches\n",
    "\n",
    "from torch_harmonics.plotting import imshow_sphere, plot_sphere\n",
    "\n",
    "from torch_harmonics.examples import Stanford2D3DSDownloader,StanfordSegmentationDataset, StanfordDepthDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset file at ../examples/data/2D3DS/stanford_2d3ds_dataset_ds16.h5 already exists. Skipping download and conversion. If you want to create a new dataset file, delete or rename the existing file.\n"
     ]
    }
   ],
   "source": [
    "data_dir = \"../examples/data/2D3DS\"\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "# 2D3DS download & dataset initialization\n",
    "downloader = Stanford2D3DSDownloader(base_url=\"https://cvg-data.inf.ethz.ch/2d3ds/no_xyz/\", local_dir=str(data_dir))\n",
    "dataset_file = downloader.prepare_dataset([(\"area_1_no_xyz.tar\", \"area_1\")], \"stanford_2d3ds_dataset_ds16.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "h5file = h5.File(os.path.join(data_dir, \"stanford_2d3ds_dataset.h5\"), \"r\")\n",
    "\n",
    "inp = np.transpose(h5file[\"rgb\"][0, ...], (1,2,0))\n",
    "tar = h5file[\"semantic\"][0, ...]\n",
    "dep = h5file[\"depth\"][0, ...]\n",
    "class_labels = h5file[\"class_labels\"]\n",
    "class_histogram = h5file[\"class_histogram\"]\n",
    "num_classes = len(class_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<HDF5 dataset \"class_labels\": shape (14,), type \"|O\">"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_labels\n",
    "#class_labels[0] = \"unknown\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7196cdef39d0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 650944 characters]",
      "text/plain": [
       "<Figure size 1800x900 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_samples = 4\n",
    "\n",
    "fig = plt.figure(layout=\"constrained\", figsize=(12, 6), dpi=150)\n",
    "subfigs = fig.subfigures(num_samples, 3, wspace=0.04)\n",
    "\n",
    "cmap = mpl.colormaps[\"tab20c\"]\n",
    "\n",
    "for isample in range(num_samples):\n",
    "\n",
    "    inp = np.transpose(h5file[\"rgb\"][isample, ...], (1,2,0))\n",
    "    tar = h5file[\"semantic\"][isample, ...]\n",
    "\n",
    "    # do the plotting\n",
    "    imshow_sphere(inp, fig=subfigs[isample, 0])\n",
    "    plot_sphere(tar / num_classes, fig =subfigs[isample, 1], cmap=cmap, vmin=0.0, vmax=1.0)\n",
    "\n",
    "    legend_list = []\n",
    "    for cidx, cname in enumerate(class_labels):\n",
    "        legend_list.append(mpatches.Patch(color=cmap(cidx / num_classes), label=cname))\n",
    "\n",
    "# Add column labels\n",
    "column_labels = [\"Input\", \"Truth\"]\n",
    "for j, label in enumerate(column_labels):\n",
    "    fig.text((j + 0.5) / 3, 1.03, label, va=\"top\", ha=\"center\", fontsize=\"large\")\n",
    "\n",
    "fig.legend(loc=\"lower center\", handles=legend_list, ncol=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<HDF5 dataset \"class_histogram\": shape (14,), type \"<f4\">"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_labels = [cl.decode(\"utf-8\") for cl in class_labels]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10340 characters]",
      "text/plain": [
       "<Figure size 864x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(12,6))\n",
    "plt.bar(class_labels, class_histogram)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 767132 characters]",
      "text/plain": [
       "<Figure size 1800x900 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_samples = 4\n",
    "\n",
    "fig = plt.figure(layout=\"constrained\", figsize=(12, 6), dpi=150)\n",
    "subfigs = fig.subfigures(num_samples, 3, wspace=0.04)\n",
    "\n",
    "cmap = mpl.colormaps[\"tab20c\"]\n",
    "\n",
    "for isample in range(num_samples):\n",
    "\n",
    "    inp = np.transpose(h5file[\"rgb\"][isample, ...], (1,2,0))\n",
    "    inp[0:20, :, :] = [1, 1, 1]\n",
    "    inp[-20:, :, :] = [1, 1, 1]\n",
    "    tar = h5file[\"depth\"][isample, ...]\n",
    "    mask = h5file[\"depth_mask\"][isample, ...]\n",
    "\n",
    "    # do the plotting\n",
    "    imshow_sphere(inp, fig=subfigs[isample, 0])\n",
    "    imshow_sphere(np.log(tar)*mask, fig =subfigs[isample, 1], cmap=\"viridis\")\n",
    "\n",
    "# Add column labels\n",
    "column_labels = [\"Input\", \"Truth\", \"Prediction\"]\n",
    "for j, label in enumerate(column_labels):\n",
    "    fig.text((j + 0.5) / 3, 1.03, label, va=\"top\", ha=\"center\", fontsize=\"large\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max:  0.99998474\n",
      "Min:  0.00077819824\n",
      "Mean log depth: 0.01618235744535923\n",
      "Std log depth: 0.011442126706242561\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7489dee04380>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 14704 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "min_depth = h5file[\"min_depth\"][...]\n",
    "max_depth = h5file[\"max_depth\"][...]\n",
    "mean_depth = h5file[\"mean_depth\"][...]\n",
    "std_depth = h5file[\"std_depth\"][...]\n",
    "\n",
    "depth = h5file[\"depth\"][...]\n",
    "mask_depth = h5file[\"depth_mask\"][...]\n",
    "\n",
    "print(\"Max: \", depth.max())\n",
    "print(\"Min: \", depth.min())\n",
    "m = np.log(1 + depth) * mask_depth\n",
    "m = m[m!=0]\n",
    "print(f\"Mean log depth: {np.nanmean(m)}\")\n",
    "print(f\"Std log depth: {np.nanstd(m)}\")\n",
    "m = (m - np.mean(m)) / np.std(m)\n",
    "\n",
    "counts, bin_edges, _ = plt.hist(m.flatten(), 1000, label=\"Masked log depth\")\n",
    "plt.legend()\n",
    "#for c,b in zip(counts, bin_edges):\n",
    "#    print(c, b)\n",
    "#print(f\"Min log depth: {np.nanmin(res)}\")\n",
    "#print(f\"Max log depth: {np.nanmax(res)}\")\n",
    "#print(f\"Mean log depth: {np.nanmean(res)}\")\n",
    "#print(f\"Std log depth: {np.nanstd(res)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Min rgb: [0. 0. 0.]\n",
      "Max rgb: [0.84705883 0.8        0.81960785]\n",
      "Mean rgb: [0.49857375 0.48026854 0.42882833]\n",
      "Std rgb: [0.38514462 0.38509178 0.38594526]\n"
     ]
    }
   ],
   "source": [
    "min_rgb = h5file[\"min_rgb\"][...]\n",
    "max_rgb = h5file[\"max_rgb\"][...]\n",
    "mean_rgb = h5file[\"mean_rgb\"][...]\n",
    "std_rgb = h5file[\"std_rgb\"][...]\n",
    "\n",
    "print(f\"Min rgb: {min_rgb}\")\n",
    "print(f\"Max rgb: {max_rgb}\")\n",
    "print(f\"Mean rgb: {mean_rgb}\")\n",
    "print(f\"Std rgb: {std_rgb}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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": 4
}
