notebooks/label_studio_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ground truth database of LabelStudio\n",
    "\n",
    "Space brings several advantages when being used as the storage of ML data labeling services. It provides simple APIs to add and remove data entries (rows), and support deduplication and overwriting data. Version management features (snapshots, tags) allow you to time travel to a previous version. Branch support is in-progress, which can modify a previous version.\n",
    "\n",
    "Space supports analysis queries on the annotations to gain insights of data. Space's transform and materialized view (MV) help you build a data pre-processing pipeline in a few lines. The source dataset can be transformed to a training ready format to feed into training frameworks. When annotations are changed (e.g., add, drop, modify annotations), you can refresh MVs to incrementally synchronize changes.\n",
    "\n",
    "This example uses [LabelStudio](https://github.com/HumanSignal/label-studio) and Space to:\n",
    "- Label objects for training image object detection models\n",
    "- Store the labeling result in Space\n",
    "- Analyze labels using SQL\n",
    "- Build a data processing pipeline with Space transform and MV, and produce [Tensorflow format](https://www.tensorflow.org/datasets/api_docs/python/tfds/features/FeaturesDict) training data.\n",
    "\n",
    "### Label object bounding boxes in images\n",
    "\n",
    "Following [Label Studio guide](https://labelstud.io/guide/install) to install and run `label-studio`, then set up the labeling interface and import the data to label.\n",
    "\n",
    "We use a mini example of [RectangleLabels](https://labelstud.io/tags/rectanglelabels) as the labeling interface:\n",
    "\n",
    "```\n",
    "<View>html\n",
    "  <Image name=\"image\" value=\"$image\"/>\n",
    "  <RectangleLabels name=\"label\" toName=\"image\">\n",
    "    <Label value=\"building\" background=\"green\"/>\n",
    "    <Label value=\"car\" background=\"red\"/>\n",
    "    <Label value=\"human\" background=\"yellow\"/>\n",
    "    <Label value=\"airplane\" background=\"gray\"/>\n",
    "    <Label value=\"boat\" background=\"blue\"/>\n",
    "  </RectangleLabels>\n",
    "</View>\n",
    "```\n",
    "\n",
    "In the next step, we import images to label to LabelStudio. This example simply uses [image files in a local directory](https://labelstud.io/guide/tasks#Import-data-from-a-local-directory). After manual labeling is done, export the annotations as JSON files. The annotations look like (pruned fields):\n",
    "\n",
    "```json\n",
    "[\n",
    "  {\n",
    "    \"id\": 8,\n",
    "    \"annotations\": [\n",
    "      {\n",
    "        \"result\": [\n",
    "          {\n",
    "            \"original_width\": 1500,\n",
    "            \"original_height\": 2060,\n",
    "            \"value\": {\n",
    "              \"x\": 0.9688418577307466,\n",
    "              \"y\": 0.7054673721340388,\n",
    "              \"width\": 97.36860670194004,\n",
    "              \"height\": 92.68077601410934,\n",
    "              \"rectanglelabels\": [\n",
    "                \"building\"\n",
    "              ]\n",
    "            },\n",
    "          }\n",
    "        ],\n",
    "        \"unique_id\": \"24dd2656-41dd-4029-9ec0-aa8ce2e4c2d0\",\n",
    "      }\n",
    "    ],\n",
    "    \"data\": {\n",
    "      \"image\": \"/segment_anything/sa_000000/sa_1.jpg\"\n",
    "    },\n",
    "    \"created_at\": \"2024-01-11T16:22:23.587074Z\"\n",
    "  }\n",
    "]\n",
    "```\n",
    "\n",
    "### Write annotations into a Space dataset\n",
    "\n",
    "Read the JSON file and convert annotations to a PyArrow table `labels`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import pyarrow as pa\n",
    "\n",
    "# The exported LabelStudio annotation file.\n",
    "exported_ls_json = \"project-2-at-2024-01-11-16-27-38c5377f.json\"\n",
    "\n",
    "# Preprocess before loading to Space.\n",
    "# Drop empty fields in JSON. It is impossible to infer types for them.\n",
    "with open(exported_ls_json) as f:\n",
    "  labels_json = json.load(f)\n",
    "\n",
    "for entry in labels_json:\n",
    "  for annotation in entry[\"annotations\"]:\n",
    "    del annotation[\"draft_created_at\"]\n",
    "    del annotation[\"prediction\"]\n",
    "    del annotation[\"import_id\"]\n",
    "    del annotation[\"last_action\"]\n",
    "    del annotation[\"parent_prediction\"]\n",
    "    del annotation[\"parent_annotation\"]\n",
    "    del annotation[\"last_created_by\"]\n",
    "  del entry[\"drafts\"]\n",
    "  del entry[\"predictions\"]\n",
    "  del entry[\"meta\"]\n",
    "  del entry[\"last_comment_updated_at\"]\n",
    "  del entry[\"comment_authors\"]\n",
    "\n",
    "# Convert the JSON array to a PyArrow table `labels`.\n",
    "labels = pa.Table.from_pandas(pd.json_normalize(labels_json))\n",
    "\n",
    "# TODO: it is a Space limitation, to remove after timestamp is supported.\n",
    "labels = labels.drop(\"created_at\")\n",
    "labels = labels.drop(\"updated_at\")\n",
    "\n",
    "# Check total number of rows.\n",
    "labels.num_rows\n",
    "\n",
    "# Check all fields.\n",
    "labels.schema.names\n",
    "# >>>\n",
    "# ['id', 'annotations', 'inner_id', 'total_annotations', 'cancelled_annotations',\n",
    "# 'total_predictions', 'comment_count', 'unresolved_comment_count', 'project', 'updated_by',\n",
    "# 'data.image']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an empty Space dataset, using the `labels` PyArrow table's schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from space import Dataset\n",
    "\n",
    "label_ds_location = \"/space/labelstudio/label_ds\"\n",
    "\n",
    "# Create an empty dataset using the `labels`'s schema.\n",
    "label_ds = Dataset.create(label_ds_location, labels.schema,\n",
    "  primary_keys=[\"id\"], record_fields=[])\n",
    "\n",
    "# Load dataset after creation:\n",
    "# label_ds = Dataset.load(label_ds_location)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data mutations and version management\n",
    "\n",
    "Append `labels` into the empty Space dataset, then delete some rows. Add a tag after each mutation, and read old data versions using tags.\n",
    "\n",
    "Branch features (work in progress) will support modifying a previous version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow.compute as pc\n",
    "\n",
    "# Append `labels` into the dataset.\n",
    "label_ds.local().append(labels)\n",
    "# Tag this version.\n",
    "label_ds.add_tag(\"after_append\")\n",
    "\n",
    "# Check all `id`s we have, and delete 2 rows.\n",
    "label_ds.local().read_all(fields=[\"id\"])\n",
    "label_ds.local().delete((pc.field(\"id\") == 8) | (pc.field(\"id\") == 9))\n",
    "label_ds.add_tag(\"after_delete\")\n",
    "\n",
    "# Read an old version.\n",
    "label_ds.local().read_all(version=\"after_append\", fields=[\"id\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Analyzing data using SQL queries\n",
    "\n",
    "The following example uses SQL queries to analyze the distribution of annotations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import duckdb\n",
    "\n",
    "# Load `annotations` column into memory.\n",
    "data = label_ds.local().read_all(fields=[\"annotations\"])\n",
    "\n",
    "# Run a SQL query to compute the max object size.\n",
    "sql = \"\"\"\n",
    "SELECT MAX(result.value.width * result.value.height) FROM (\n",
    "  SELECT UNNEST(annotations.result) AS result FROM (\n",
    "    SELECT UNNEST(annotations) AS annotations FROM data\n",
    "  )\n",
    ")\n",
    "\"\"\"\n",
    "duckdb.sql(sql).fetchall()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transform annotations to a training ready Tensorflow dataset\n",
    "\n",
    "[Tensorflow dataset](https://www.tensorflow.org/datasets/tfless_tfds) stores training data in [ArrayRecord](https://github.com/google/array_record) files. It serializes data into bytes using [FeaturesDict](https://www.tensorflow.org/datasets/api_docs/python/tfds/features/FeaturesDict). An ArrayRecord file stores an array of bytes (serialized features).\n",
    "\n",
    "We define a simple `FeaturesDict` for training an object detector:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tensorflow_datasets import features as f\n",
    "\n",
    "tf_features_dict = f.FeaturesDict({\n",
    " \"image\": f.Image(shape=(None, None, 3), dtype=np.uint8),\n",
    " \"objects\": f.Sequence({\n",
    "   \"bbox\": f.BBoxFeature(),\n",
    "   \"label\": f.ClassLabel(num_classes=5),\n",
    "  }),\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`preprocess_labels` is a user defined function whose input is a batch of `id`, `data.image`, `annotations` columns of the dataset `label_ds`. It reads and resizes the image, and converts bounding boxes to the TFDS format. At the end, it serializes the features and returns the converted batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict\n",
    "import cv2\n",
    "\n",
    "label_classes = {\n",
    "  \"airplane\": 0,\n",
    "  \"boat\": 1,\n",
    "  \"building\": 2,\n",
    "  \"human\": 3,\n",
    "  \"car\": 4\n",
    "}\n",
    "\n",
    "def preprocess_labels(data: Dict[str, Any]) -> Dict[str, Any]:\n",
    "  features = []\n",
    "  for file_path, annotations in zip(data[\"data.image\"], data[\"annotations\"]):\n",
    "    # NOTE: may need to modify file_path in your case.\n",
    "    im = cv2.imread(file_path)\n",
    "    im = cv2.resize(im, dsize=(100, 100), interpolation=cv2.INTER_CUBIC)\n",
    "\n",
    "    objects = []\n",
    "    for annotation in annotations:\n",
    "      for result in annotation[\"result\"]:\n",
    "        v = result[\"value\"]\n",
    "        ymin, xmin = v[\"y\"] / 100, v[\"x\"] / 100\n",
    "        w, h = v[\"width\"] / 100, v[\"height\"] / 100\n",
    "        ymax, xmax = ymin + h, xmin + w\n",
    "        objects.append({\n",
    "          \"bbox\": f.BBox(ymin, xmin, ymax, xmax),\n",
    "          \"label\": label_classes[v[\"rectanglelabels\"][0]]\n",
    "        })\n",
    "\n",
    "    # Serialize the TFDS feature to bytes to write to storage.\n",
    "    features.append(\n",
    "      tf_features_dict.serialize_example({\"image\": im, \"objects\": objects}))\n",
    "\n",
    "  return {\"id\": data[\"id\"], \"features\": features}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use a `map` transform to convert the dataset `label_ds` to a view `training_view`. Then create a materialized view. When `label_ds` is modified, simply `refresh` the MV to incrementally synchronize changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from space import TfFeatures \n",
    "\n",
    "training_view = label_ds.map_batches(\n",
    "  fn=preprocess_labels,\n",
    "  input_fields=[\"id\", \"data.image\", \"annotations\"],  # Input of fn\n",
    "  output_schema=pa.schema([   # Schema of the view\n",
    "    (\"id\", pa.int64()),\n",
    "    (\"features\", TfFeatures(tf_features_dict))\n",
    "  ]),\n",
    "  output_record_fields=[\"features\"]  # Store this field in ArrayRecord\n",
    ")\n",
    "\n",
    "training_mv_location = \"/space/labelstudio/training_mv\"\n",
    "training_mv = training_view.materialize(training_mv_location)\n",
    "\n",
    "training_mv.ray().refresh(\"after_append\")\n",
    "\n",
    "# from space import MaterializedView\n",
    "# training_mv = MaterializedView.load(training_mv_location)\n",
    "# training_mv.ray().read_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use Space's `RandomAccessDataSource` to feed data to the training framework:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from space import RandomAccessDataSource\n",
    "\n",
    "datasource = RandomAccessDataSource(\n",
    "  # field-name: storage-location, for reading data from ArrayRecord files.\n",
    "  {\"features\": training_mv_location},\n",
    "  # Auto deserialize data using `tf_features_dict`.\n",
    "  deserialize=True)\n",
    "\n",
    "# Read the data source.\n",
    "len(datasource)\n",
    "\n",
    "datasource[1]\n",
    "# >>>\n",
    "# {'image': array([[[234, 227, 224],\n",
    "#         [243, 237, 238],\n",
    "#         [245, 239, 236],\n",
    "#         ...,\n",
    "#         [187, 176, 178]]], dtype=uint8),\n",
    "#  'objects': {'bbox': array([[0.00705467, 0.00968842, 0.93386245, 0.9833745 ]], dtype=float32),\n",
    "#              'label': array([2])}}"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}