tutorials/load_and_sample_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bca11a49-551d-49bb-970b-460313d13c22",
   "metadata": {},
   "source": [
    "# Tutorial: Load and Sample Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6f4c13d-af8e-4dd4-bdc0-8701922cc757",
   "metadata": {},
   "source": [
    "In this tutorial you will learn what are the different ways to load datasets, subsample them.\n",
    "\n",
    "To run all cells, you need the `annotated-trainval` and `synthetic` datasets unzipped in `data/docile` and `annotated-trainval.zip` file in `data/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf865fb-be3d-4243-b261-7b6b8c46bf6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from docile.dataset import CachingConfig, Dataset\n",
    "\n",
    "DATASET_PATH = Path(\"/app/data/docile/\")\n",
    "DATASET_PATH_ZIP = Path(\"/app/data/annotated-trainval.zip\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e17c01c-9d6f-48d6-9bff-02f73244bdfd",
   "metadata": {},
   "source": [
    "## Load from folder with unzipped dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93784f80-98e2-4375-9df0-3e0dcec0fd51",
   "metadata": {},
   "outputs": [],
   "source": [
    "val = Dataset(\"val\", DATASET_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b26393b-dd4a-42e2-87fb-7fc154d1efab",
   "metadata": {},
   "source": [
    "## Load from zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f82f329-f545-4303-a4de-720ef8ecd66e",
   "metadata": {},
   "source": [
    "Dataset can be loaded directly from zip as well but image caching to disk must be turned off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a09ffdb3-93bc-4ffc-9b13-d4770157fa73",
   "metadata": {},
   "outputs": [],
   "source": [
    "val = Dataset(\"val\", DATASET_PATH_ZIP, cache_images=CachingConfig.OFF)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f77805d5-513e-45b3-aa29-d721a8315a76",
   "metadata": {},
   "source": [
    "## Preloading document resources to memory and image caching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bca6070d-59f4-41e4-976c-a69262638e24",
   "metadata": {},
   "source": [
    "By default, dataset is loaded with these settings:\n",
    "\n",
    "* annotations and pre-computed OCR are loaded to memory\n",
    "* images generated from PDFs are cached to disk (for faster access of the images in future iterations)\n",
    "\n",
    "Below you see options how to change this default behaviour, which is especially useful for large datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ad98f3-ca23-4abb-8d40-075c396bed2a",
   "metadata": {},
   "source": [
    "**Only load annotations, not pre-computed OCR**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bdbdca9-5c01-4975-ab6f-850e333a2026",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = Dataset(\"train\", DATASET_PATH, load_ocr=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ce69236-b160-479b-8d34-892f30550d72",
   "metadata": {},
   "source": [
    "**Postpone loading of document resources**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a1bdfe1-c42d-4d06-8566-b15bdec8ab65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do not preload document resources\n",
    "synthetic = Dataset(\"synthetic\", DATASET_PATH, load_annotations=False, load_ocr=False, cache_images=CachingConfig.OFF)\n",
    "# You can load part of the dataset later\n",
    "synthetic_sample = synthetic[:100].load()\n",
    "# And release it from memory later\n",
    "synthetic_sample.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0725f950-4cf3-427a-a6b6-1bedda063f95",
   "metadata": {},
   "source": [
    "**Cache images to both disk and memory**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ee58b4-8b6d-4a9e-9f9b-858a0403be0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cache images also in memory. Make sure you have enough RAM memory to do this!\n",
    "# Images are not loaded to memory right away but only after first\n",
    "train = Dataset(\"train\", DATASET_PATH, cache_images=CachingConfig.DISK_AND_MEMORY)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d8456a0-7a3c-42f0-b368-dfb1bbbc6b29",
   "metadata": {},
   "source": [
    "## Sample and chunk documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb21a02-da04-4aa8-9ae0-d430f9fc9a82",
   "metadata": {},
   "source": [
    "For experiments and to work with large datasets, it can be useful to take samples of the datasets.\n",
    "\n",
    "For this, you can use slicing `[start:end:step]`, `.sample()`, `.get_cluster()` or `.from_documents()` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a69741c7-4b10-4409-86e9-2986a59b3057",
   "metadata": {},
   "outputs": [],
   "source": [
    "synthetic = Dataset(\"synthetic\", DATASET_PATH, load_annotations=False, load_ocr=False, cache_images=CachingConfig.OFF)\n",
    "trainval = Dataset(\"trainval\", DATASET_PATH, load_annotations=False, load_ocr=False, cache_images=CachingConfig.OFF)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9621d77e-5820-4b2b-967f-98771468c07d",
   "metadata": {},
   "source": [
    "**Slicing**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c9cab7d-ac31-4cd3-a3c0-07ae91001cb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Synthetic document has 100 chunks of 1000 documents from the same template document, so the\n",
    "# following line selects 1 document for each template document:\n",
    "synthetic_slice = synthetic[8::1000]\n",
    "print(synthetic_slice.docids[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "788090cc-d536-4cc3-8df3-9e43acd4d598",
   "metadata": {},
   "source": [
    "**Random sample**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20041b46-49f5-475c-9fcb-ff2dea15ba2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "synthetic_sample = synthetic.sample(5)\n",
    "print(synthetic_sample)\n",
    "print(synthetic_sample.docids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "263df176-3366-4a25-8ee3-34d722ef0c79",
   "metadata": {},
   "source": [
    "**Documents belonging to the same cluster**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3026aabd-c303-4e0a-aced-4d6501b8a4d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainval_cluster = trainval.get_cluster(synthetic[0].annotation.cluster_id)\n",
    "print(f\"Found {len(trainval_cluster)} documents in {trainval_cluster}.\")\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "print(\"Showing 10 images from the cluster:\")\n",
    "imgs = [doc.page_image(page=0, image_size=(None, 100)) for doc in trainval_cluster[:10]]\n",
    "concat_img = Image.new(\"RGB\", (sum(img.width for img in imgs), 100))\n",
    "start_from = 0\n",
    "for img in imgs:\n",
    "    concat_img.paste(img, (start_from, 0))\n",
    "    start_from += img.width\n",
    "concat_img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6928e572-09c7-4637-b2a2-40f2e2967950",
   "metadata": {},
   "source": [
    "**Using custom filter**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0724681f-9a13-4a8a-bad9-5c10b3473a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainval_ucsf = Dataset.from_documents(\"trainval-ucsf\", [doc for doc in trainval if doc.annotation.source == \"ucsf\"])\n",
    "trainval_pif = Dataset.from_documents(\"trainval-pif\", [doc for doc in trainval if doc.annotation.source == \"pif\"])\n",
    "print(f\"{trainval_ucsf} with {len(trainval_ucsf)} documents\")\n",
    "print(f\"{trainval_pif} with {len(trainval_pif)} documents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0999723a-e164-4bf4-9409-bcc68a458989",
   "metadata": {},
   "source": [
    "## Chunk dataset into parts with the same number of pages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf55f32c-1fc9-4291-be19-c1c3d6a384e6",
   "metadata": {},
   "source": [
    "Create dataset chunks that have a limited number of pages. This can be especially useful for large datasets, such as the unlabeled dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bed2781-b19e-4ce5-9270-0c000a2647a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Iterable\n",
    "\n",
    "def chunk_dataset(dataset: Dataset, max_pages_per_chunk: int) -> Iterable[Dataset]:\n",
    "    start_doc_i = 0\n",
    "    pages = 0\n",
    "    for doc_i, document in enumerate(dataset.documents):\n",
    "        documents = doc_i - start_doc_i + 1\n",
    "        pages += document.page_count\n",
    "        if doc_i > start_doc_i and pages > max_pages_per_chunk:\n",
    "            yield dataset[start_doc_i:doc_i]\n",
    "            start_doc_i = doc_i\n",
    "            pages = document.page_count\n",
    "    yield dataset[start_doc_i:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3645bee-c173-49aa-b301-2b4e1fd49b21",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainval = Dataset(\"trainval\", DATASET_PATH, load_annotations=False, load_ocr=False, cache_images=CachingConfig.OFF)\n",
    "\n",
    "max_pages_per_chunk = 2000\n",
    "\n",
    "for chunk in chunk_dataset(trainval, max_pages_per_chunk):\n",
    "    print(f\"{chunk}, pages: {chunk.total_page_count()}\")\n",
    "    chunk.load(annotations=True, ocr=False)\n",
    "    # ... work with the chunk here ...\n",
    "    chunk.release() # don't forget to free up the memory"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}