{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f163b281",
   "metadata": {},
   "source": [
    "# Vector Database Demo\n",
    "\n",
    "A vector database allows for efficient similarity searches by managing a data structure on behalf of the user that defines a distant between any two vectors. This allows the user to query the database for similar vectors given an input vector. In this notebook, we'll see how to use the Hyrax to work with the built in vector database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f174ad0",
   "metadata": {},
   "source": [
    "As usual, we'll start by creating an instance of a `Hyrax` object and then quickly `train` the default HyraxAutoencoder. After than we'll run some data through `infer` to produce our encoded vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "44b00d84",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 14:59:59,703 hyrax:INFO] Runtime Config read from: /home/drew/code/hyrax/src/hyrax/hyrax_default_config.toml\n"
     ]
    }
   ],
   "source": [
    "from hyrax import Hyrax\n",
    "\n",
    "h = Hyrax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e214de6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 15:00:12,510 hyrax.models.model_registry:INFO] Using criterion: torch.nn.CrossEntropyLoss with default arguments.\n",
      "2025-05-20 15:00:12,568 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': <hyrax.pytorch_ignite.SubsetSequentialSampler object at 0x7fccbe457350>, 'batch_size': 512, 'shuffle': False, 'pin_memory': True}\n",
      "2025-05-20 15:00:12,569 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': <hyrax.pytorch_ignite.SubsetSequentialSampler object at 0x7fccbe1b8440>, 'batch_size': 512, 'shuffle': False, 'pin_memory': True}\n",
      "/home/drew/miniconda3/envs/hyrax/lib/python3.12/site-packages/ignite/handlers/tqdm_logger.py:127: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n",
      "2025/05/20 15:00:12 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.\n",
      "[2025-05-20 15:00:12,823 hyrax.pytorch_ignite:INFO] Training model on device: cuda\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd6cf99391d5409f916e2e53981ffafd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "626b046abf824ba584177ea938df3c2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc513e62417b492697c4132c3b6663c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07a62bf384d14aa8bc8a2feb18679cae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53cc4958796447539bfded2bc5878348",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "954d5d3175614c58b8ce128aef1fcf66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e9962399f4f34948b0e6dcec9edfb87d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9f54959d01e46629b711800b04da14d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "477d1073fa2a43d7976f9f1f9e3a7bd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f23f7b382a0e45e0ab352fc1475f6238",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aacf976f2ada44818cdfbdbed2e65d95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "254a435980c04c9a92396482f2e8ae56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50ae71d9dd08476881b7abddeddacfea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a1b6067537941f7b1d34159257c55cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3f58a3f5ca7a4a19baaee0012339843a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad52b67aef834936af8ac3e049524c02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbd46f47f8fd4edf91e22887dc5b092d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7605c7b2cf434cce875453928aea5ced",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3f092339ffeb46c3b76b68272f991880",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0c2f1f693dc3433ea8f26d1aed59e5b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04200b31d6704bdcb6054df85a62e470",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4413109d28af47fd956e969c597a293a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44c2daa1f2134b078d1a43c5380e818c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5a8015acc13426d8b717c902fc82663",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96c3d9913fbb4a4899375ebe7d0c0a42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 15:04:14,099 hyrax.pytorch_ignite:INFO] Total training time: 241.28[s]\n",
      "[2025-05-20 15:04:14,100 hyrax.pytorch_ignite:INFO] Latest checkpoint saved as: /home/drew/code/hyrax/docs/pre_executed/results/20250520-150004-train-Is49/checkpoint_epoch_25.pt\n",
      "[2025-05-20 15:04:14,100 hyrax.pytorch_ignite:INFO] Best metric checkpoint saved as: /home/drew/code/hyrax/docs/pre_executed/results/20250520-150004-train-Is49/checkpoint_25_loss=-91.8765.pt\n",
      "2025/05/20 15:04:14 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...\n",
      "2025/05/20 15:04:14 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!\n",
      "[2025-05-20 15:04:14,115 hyrax.train:INFO] Finished Training\n",
      "[2025-05-20 15:04:14,391 hyrax.model_exporters:INFO] Exported model to ONNX format: /home/drew/code/hyrax/docs/pre_executed/results/20250520-150004-train-Is49/example_model_opset_20.onnx\n",
      "[2025-05-20 15:04:22,593 hyrax.models.model_registry:INFO] Using criterion: torch.nn.CrossEntropyLoss with default arguments.\n",
      "[2025-05-20 15:04:22,593 hyrax.verbs.infer:INFO] data set has length 50000\n",
      "2025-05-20 15:04:22,595 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': None, 'batch_size': 512, 'shuffle': False, 'pin_memory': True}\n",
      "[2025-05-20 15:04:22,618 hyrax.verbs.infer:INFO] Saving inference results at: /home/drew/code/hyrax/docs/pre_executed/results/20250520-150414-infer-v7i0\n",
      "[2025-05-20 15:04:22,887 hyrax.pytorch_ignite:INFO] Evaluating model on device: cuda\n",
      "[2025-05-20 15:04:22,889 hyrax.pytorch_ignite:INFO] Total epochs: 1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "25594032566344b29a35719f2107b185",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  1%|1         | 1/98 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 15:04:33,962 hyrax.pytorch_ignite:INFO] Total evaluation time: 11.07[s]\n",
      "[2025-05-20 15:04:34,088 hyrax.verbs.infer:INFO] Inference Complete.\n",
      "[2025-05-20 15:04:34,143 hyrax.config_utils:WARNING] Cannot find default_config.toml for torch.\n",
      "[2025-05-20 15:04:34,144 hyrax.config_utils:WARNING] Cannot find default_config.toml for torch.\n",
      "[2025-05-20 15:04:37,769 hyrax.config_utils:WARNING] Cannot find default_config.toml for umap.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<hyrax.data_sets.inference_dataset.InferenceDataSet at 0x7fcccb544140>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h.config[\"train\"][\"epochs\"] = 25\n",
    "h.train()\n",
    "h.infer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f2da27c",
   "metadata": {},
   "source": [
    "Now we'll use the output of inference to populate a ChromaDB vector database. By default `h.save_to_database()` will find the most recent output from `h.infer()` to use as the inputs when filling the database, but the user can specify a different dataset. Likewise, the default behavior is to create a new directory to store the vector database, but this can also be specified.\n",
    "\n",
    "`h.save_to_database([input_dir=\"/path/to/input\",] [output_dir=\"/path/to/output\"])`\n",
    "\n",
    "Note: If the goal is to add more data to an existing database, provide the existing database directory as the `output_dir` when calling `h.save_to_database`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6b8e25f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filling a chromadb database with the results of inference.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 15:04:46,286 hyrax.config_utils:WARNING] Cannot find default_config.toml for torch.\n",
      "[2025-05-20 15:04:46,287 hyrax.config_utils:WARNING] Cannot find default_config.toml for torch.\n",
      "[2025-05-20 15:04:46,288 hyrax.config_utils:WARNING] Cannot find default_config.toml for umap.\n",
      "[2025-05-20 15:04:54,737 hyrax.verbs.save_to_database:INFO] Number of inference result batches to index: 98.\n",
      "100%|██████████| 98/98 [00:12<00:00,  8.06it/s]\n"
     ]
    }
   ],
   "source": [
    "print(f\"Filling a {h.config['vector_db']['name']} database with the results of inference.\")\n",
    "h.save_to_database()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe58a5ef",
   "metadata": {},
   "source": [
    "Now that the vector database has been populated, it can be used for similarity search and nearest neighbors. First we'll need to establish a connection using `h.database_connection()`.\n",
    "\n",
    "With the connection established, we have access to three types of queries:\n",
    "* `get_by_id` to retrieve the vector associated with a particular id\n",
    "* `search_by_id` to retrieve the _k_ nearest neighbors to a particular id\n",
    "* `search_by_vector` to retrieve the _k_ nearest neighbors to a particular input vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6d22029f",
   "metadata": {},
   "outputs": [],
   "source": [
    "conn = h.database_connection()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a1934c7",
   "metadata": {},
   "source": [
    "### `get_by_id`\n",
    "\n",
    "With `get_by_id` we can retrieve the vectors associated with some ids. The result is a dictionary where the keys are id strings and values are the numpy array vectors stored in the vector database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f61c605",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Id: 334, Vector: [ 0.21406414  4.38409519  0.49727368 -0.7893768  -2.87350202  0.11476225\n",
      "  1.8274399  -0.09784815 -1.2193296  -0.77086699  1.34052861  1.65932298\n",
      " -2.07148337 -2.34664512  2.28803968  0.86179125  0.86639947 -2.3657589\n",
      " -1.03186095  1.59934366 -0.38279176 -0.92782491 -1.26808596  1.49409449\n",
      " -2.36837769 -2.97786927 -1.81079423 -1.38344455 -1.11257839  1.33133769\n",
      "  0.86067116 -2.94752669  0.02824669  1.38922858 -0.20957161 -0.15272772\n",
      "  0.01273713 -0.05502158 -0.28467956 -0.06139451 -1.35991454 -0.46939212\n",
      " -1.15573668 -3.90672135  1.14438629  4.01489925  0.87542021 -1.55964196\n",
      " -0.27592164  1.35570014 -0.8417064   0.40743572 -0.45822605  0.36299995\n",
      " -4.36139441 -0.92870319  1.6725179  -0.35636529  0.2801117  -1.55952835\n",
      " -2.01857901  0.28982323 -2.65234852  2.10563254]\n",
      "Id: 19022, Vector: [ 1.01997221  3.27541471  0.07880742  0.40190896 -2.069767    0.22248952\n",
      "  0.63698471  0.91667783 -1.42546344 -0.78596801  1.58609724  0.09808031\n",
      " -2.03873038 -1.99726641  1.99178457  0.30669096  1.55701888 -1.48524356\n",
      " -0.9965176   0.82624108 -0.4912349  -0.61169094 -0.96432102  2.02064347\n",
      " -1.09434426 -1.66586471 -0.8439483  -0.19334477 -0.20327687  1.13053894\n",
      "  0.90265465 -1.52364016  0.38151348  0.74716371 -0.0561714  -0.10911166\n",
      " -0.34963098 -0.36962882  0.66240174  0.64269614 -0.68175811 -1.75478578\n",
      " -1.9701103  -3.3509984   2.56497812  4.5865469   0.91742259 -0.29954159\n",
      " -0.07238744  1.71732593 -0.27246445 -0.20318703 -1.55732191  1.46810496\n",
      " -2.89273024 -0.77957153  1.27631938 -0.19323938  0.73189634 -1.65531635\n",
      " -1.78045166 -0.07765396 -3.09991169  1.34831798]\n"
     ]
    }
   ],
   "source": [
    "results = conn.get_by_id(ids=[\"334\", \"19022\"])\n",
    "\n",
    "for k, v in results.items():\n",
    "    print(f\"Id: {k}, Vector: {v}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14cff8dc",
   "metadata": {},
   "source": [
    "### `search_by_id`\n",
    "\n",
    "Now we'll search for nearest neighbors. First we search for the _k_ nearest neighbors of using an id from the database. Note that the closest of the 5 neighbors is the vector itself, since it's in the database. The returned dictionary contains the ids of the nearest neighbors in order of increasing distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b7d4b69e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'19022': ['19022', '431', '30105', '4532', '27167']}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conn.search_by_id(\"19022\", k=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8079aa88",
   "metadata": {},
   "source": [
    "### `search_by_vector`\n",
    "\n",
    "We can repeat the search this time using the vector that we found when running `get_by_id`. The results from this search should match the previous results from `search_by_id`. Again, the returned dictionary contains the ids of the nearest neighbors in order of increasing distance.\n",
    "\n",
    "Note that the input vector to `search_by_vector` is a list of vectors. Also the input vectors don't need to be vectors extracted from the database. For example - if a new piece of data was run through `infer`, the resulting vector could be passed as input to `search_by_vector` to find similar data. This is why the keys of the returned dictionary are integers, they are the indexes of the input vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "81088de0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: ['19022', '431', '30105', '4532', '27167']}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "neighbors = conn.search_by_vector([results[\"19022\"]], k=5)\n",
    "neighbors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a32ea79",
   "metadata": {},
   "source": [
    "We have now successfully queried our vector database. Let's take a look at the nearest neighbors that were found for the example image. First we'll get an interactive copy of the original dataset, and write a couple of helpful plotting functions so that we can plot the original images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "32946b37",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-20 15:05:15,134 hyrax.prepare:INFO] Finished Prepare\n"
     ]
    }
   ],
   "source": [
    "dataset = h.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44c22fce",
   "metadata": {},
   "source": [
    "The following is some boilerplate code for displaying either a single image or a collection of images from `dataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2d245a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "\n",
    "def show_image(data):\n",
    "    image = data[\"image\"]\n",
    "    title = f\"Label: {data['label']}, Id: {data['object_id']}\"\n",
    "    image = image.permute(1, 2, 0).numpy()\n",
    "    min_val = np.min(image)\n",
    "    max_val = np.max(image)\n",
    "    image = (image - min_val) / (max_val - min_val)\n",
    "    plt.imshow(image)\n",
    "    if title:\n",
    "        plt.title(title)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def show_image_grid(data_list):\n",
    "    \"\"\"\n",
    "    Display a grid of images.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    data_list : list\n",
    "        A list of data dictionaries, each containing an 'image' key and optionally 'label' and 'object_id'.\n",
    "    \"\"\"\n",
    "    rows = int(len(data_list) / 4)\n",
    "    if len(data_list) % 4 != 0:\n",
    "        rows += 1\n",
    "    cols = 4\n",
    "    if rows == 1:\n",
    "        cols = len(data_list)\n",
    "    fig, axes = plt.subplots(rows, cols, figsize=(cols * 3, rows * 3))\n",
    "    axes = axes.flatten()  # Flatten the axes for easy iteration\n",
    "\n",
    "    for ax, data in zip(axes, data_list):\n",
    "        image = data[\"image\"]\n",
    "        title = f\"Label: {data.get('label', 'N/A')}, Id: {data.get('object_id', 'N/A')}\"\n",
    "        image = image.permute(1, 2, 0).numpy()\n",
    "        min_val = np.min(image)\n",
    "        max_val = np.max(image)\n",
    "        image = (image - min_val) / (max_val - min_val)  # Normalize to [0, 1]\n",
    "        ax.imshow(image)\n",
    "        ax.set_title(title)\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "    # Turn off any unused axes\n",
    "    for ax in axes[len(data_list) :]:\n",
    "        ax.axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "482e3355",
   "metadata": {},
   "source": [
    "First we'll display the image that we used for searching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bd6f001f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12800 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_image(dataset[19022])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64aa055b",
   "metadata": {},
   "source": [
    "Next we'll display the 4 nearest neighbors. We'll exclude the absolute nearest neighbor, because in our case, that is the original search image. It seems that the model is doing something vaguely reasonable. By eye, some of the nearest neighbors appear similar to the original search image. Though, clearly, there is room for model improvement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8ac4f559",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33216 characters]",
      "text/plain": [
       "<Figure size 1200x300 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_list = [dataset[int(n)] for n in neighbors[0][1:]]\n",
    "show_image_grid(data_list)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hyrax",
   "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.12.10"
  },
  "nbsphinx": {
   "orphan": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
