notebooks/egoexo/EgoExo_Atomic_Descriptions_Tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c6e2450",
   "metadata": {},
   "source": [
    "# Atomic Actions Descriptions Tutorial\n",
    "\n",
    "In this tutorial, we will:\n",
    "1. Read the annotation file\n",
    "2. Compute aggregate stats on the atomic action descriptions\n",
    "3. Embed the descriptions with Sentence Transformers (https://www.sbert.net/)\n",
    "4. Perform clustering on these embedddings to find common phrases using sklearn\n",
    "5. Visualize the embeddings with Dimension Reduction (t-SNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0414f48-092e-4c5d-8b3e-353e9b551312",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import time\n",
    "import os\n",
    "import traceback\n",
    "import random\n",
    "from collections import defaultdict\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fab38a53-9a1b-4c46-9b6c-3a64a799587a",
   "metadata": {},
   "outputs": [],
   "source": [
    "RELEASE_DIR = \"/path/placeholder\"  # NOTE: changeme\n",
    "# RELEASE_DIR = \"/large_experiments/egoexo/v2/\"  # NOTE: changeme\n",
    "assert os.path.exists(RELEASE_DIR), \"change RELEASE_DIR to where you downloaded the dataset to\"\n",
    "\n",
    "egoexo = {\n",
    "    \"takes\": os.path.join(RELEASE_DIR, \"takes.json\"),\n",
    "    \"captures\": os.path.join(RELEASE_DIR, \"captures.json\"),\n",
    "    \"physical_setting\": os.path.join(RELEASE_DIR, \"physical_setting.json\"),\n",
    "    \"participants\": os.path.join(RELEASE_DIR, \"participants.json\"),\n",
    "    \"visual_objects\": os.path.join(RELEASE_DIR, \"visual_objects.json\"),\n",
    "}\n",
    "\n",
    "TASK_ID_CAT = {\n",
    "    0: \"Unknown\",\n",
    "    1000: \"Cooking\",\n",
    "    2000: \"Health\",\n",
    "    4000: \"Bike Repair\",\n",
    "    5000: \"Music\",\n",
    "    6000: \"Basketball\",\n",
    "    7000: \"Rock Climbing\",\n",
    "    8000: \"Soccer\",\n",
    "    9000: \"Dance\",\n",
    "}\n",
    "\n",
    "for k, v in egoexo.items():\n",
    "    egoexo[k] = json.load(open(v))\n",
    "\n",
    "takes = egoexo[\"takes\"]\n",
    "captures = egoexo[\"captures\"]\n",
    "takes_by_uid = {x[\"take_uid\"]: x for x in takes}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda172bc-6a05-45c8-bf73-7e8c50d4e96e",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation_dir = os.path.join(RELEASE_DIR, \"annotations/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93af2ec-2edf-4bb1-bcb3-b06c387c604b",
   "metadata": {},
   "outputs": [],
   "source": [
    "atomic_anns = json.load(open(os.path.join(annotation_dir, \"atomic_descriptions_train.json\")))\n",
    "anns = atomic_anns[\"annotations\"]\n",
    "print(len(anns))\n",
    "\n",
    "old_len = len(anns)\n",
    "anns.update(json.load(open(os.path.join(annotation_dir, \"atomic_descriptions_val.json\")))[\"annotations\"])\n",
    "print(len(anns), len(anns) - old_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f87404b7",
   "metadata": {},
   "source": [
    "## Read the Annotation File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1270fd02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# as a flat list\n",
    "all_descs = [\n",
    "    (take_uid, x)\n",
    "    for take_uid, xs in anns.items()\n",
    "    for y in xs\n",
    "    for x in y[\"descriptions\"]\n",
    "]\n",
    "all_descs[0:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acd39f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_anns = [y for xs in anns.values() for y in xs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26cd7eb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grouped by the task the take is associated to.\n",
    "# Below also retrieves (for stats):\n",
    "# 1. the coverage per task\n",
    "# 2. the density of annotations (relating to take time)\n",
    "coverage_by_task = defaultdict(lambda: defaultdict(int))\n",
    "desc_per_minute = defaultdict(list)\n",
    "desc_by_task = defaultdict(list)\n",
    "for take_uid, vs in anns.items():\n",
    "    take = takes_by_uid.get(take_uid)\n",
    "    if take is None:\n",
    "        task = \"Redacted\"\n",
    "        continue\n",
    "    else:\n",
    "        task = int(take[\"task_id\"]) // 1000 if take[\"task_id\"] is not None else \"Dropped\"\n",
    "\n",
    "    cat_name = task\n",
    "    if isinstance(task, int):\n",
    "        cat_name = TASK_ID_CAT[task * 1000]\n",
    "    task = cat_name\n",
    "    \n",
    "    # get all the annotated descriptions by their task\n",
    "    for x in vs:\n",
    "        desc_by_task[task].extend(x[\"descriptions\"])\n",
    "        desc_by_task[\"all\"].extend(x[\"descriptions\"])\n",
    "    \n",
    "    # compute coverage\n",
    "    coverage_by_task[task][len(vs)] += 1\n",
    "    coverage_by_task[\"all\"][len(vs)] += 1\n",
    "    \n",
    "    # compute density\n",
    "    if take is not None:\n",
    "        for y in vs:\n",
    "            num_narr = len(y[\"descriptions\"])\n",
    "            dur = (take[\"duration_sec\"] / 60)\n",
    "            narrs_per_min = num_narr / dur\n",
    "            # desc_per_minute[task].append((take_uid, narrs_per_min, take[\"duration_sec\"], num_narr, dur))\n",
    "            desc_per_minute[task].append(narrs_per_min)\n",
    "            desc_per_minute[\"all\"].append(narrs_per_min)\n",
    "\n",
    "dict(coverage_by_task[\"all\"]), coverage_by_task.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa4e585",
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_by_task[\"all\"][0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff6b6d2",
   "metadata": {},
   "source": [
    "## Compute & Display Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb88cc4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "HAS_SPACY = False\n",
    "PREFER_SPACY = False  # if you want stats on nouns/verbs\n",
    "\n",
    "if PREFER_SPACY:\n",
    "    try:\n",
    "        import spacy\n",
    "        nlp = spacy.load(\"en_core_web_md\")\n",
    "        print(\"spacy loaded!\", flush=True)\n",
    "        HAS_SPACY = True\n",
    "    except Exception:\n",
    "        print(f\"WARN: spacy could not be loaded. This is not necessary to run the notebook. :\\n{traceback.format_exc()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "728b9611",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_desc_stats():\n",
    "    return {\n",
    "        \"num_nouns\": [],\n",
    "        \"num_verbs\": [],\n",
    "        \"num_sents\": [],\n",
    "        \"num_words\": [],\n",
    "        \"words_per_sentence\": [],\n",
    "    }\n",
    "\n",
    "# get the unique nouns & verbs by category\n",
    "noun_counts = defaultdict(lambda: defaultdict(int))\n",
    "verb_counts = defaultdict(lambda: defaultdict(int))\n",
    "desc_stats = defaultdict(create_desc_stats)\n",
    "\n",
    "if HAS_SPACY:\n",
    "    print(\"Processing with spacy. WARN: this will take ~15 minutes\")\n",
    "    for take_uid, x in tqdm(all_descs):\n",
    "        take = takes_by_uid.get(take_uid)\n",
    "        if take is None:\n",
    "            task = \"Redacted\"\n",
    "            continue\n",
    "        else:\n",
    "            task = int(take[\"task_id\"]) // 1000 if take[\"task_id\"] is not None else \"Dropped\"\n",
    "\n",
    "        doc = nlp(x[\"text\"])\n",
    "        num_sents = len(list(doc.sents))\n",
    "        num_words = len(doc)\n",
    "        words_per_sentence = num_words / num_sents if num_sents > 0 else None\n",
    "\n",
    "        # group the tokens by their class (noun, verb, etc.)\n",
    "        toks_by_class = defaultdict(list)\n",
    "        for tok in doc:\n",
    "            toks_by_class[tok.pos_].append(tok)\n",
    "\n",
    "        num_nouns = len(toks_by_class[\"NOUN\"]) + len(toks_by_class[\"PROPN\"])\n",
    "        num_verbs = len(toks_by_class[\"VERBS\"])\n",
    "        for tok in toks_by_class[\"NOUN\"]:\n",
    "            noun_counts[\"all\"][tok.text] += 1\n",
    "            noun_counts[task][tok.text] += 1\n",
    "        for tok in toks_by_class[\"PROPN\"]:\n",
    "            noun_counts[\"all\"][tok.text] += 1\n",
    "            noun_counts[task][tok.text] += 1\n",
    "        for tok in toks_by_class[\"VERB\"]:\n",
    "            if tok.text == \"'s\":\n",
    "                continue\n",
    "            verb_counts[\"all\"][tok.text] += 1\n",
    "            verb_counts[task][tok.text] += 1\n",
    "\n",
    "        for part in [task, \"all\"]:\n",
    "            desc_stats[part][\"num_nouns\"].append(num_nouns)\n",
    "            desc_stats[part][\"num_verbs\"].append(num_verbs)\n",
    "            desc_stats[part][\"num_sents\"].append(num_sents)\n",
    "            desc_stats[part][\"num_words\"].append(num_words)\n",
    "            desc_stats[part][\"words_per_sentence\"].append(words_per_sentence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e478667",
   "metadata": {},
   "outputs": [],
   "source": [
    "desc_stats_df = pd.DataFrame(desc_stats[\"all\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc8a881",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_descs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24a19cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_takes_covered = len(anns)\n",
    "desc_per_ann = np.array([len(y[\"descriptions\"]) for y in all_anns])\n",
    "noun_counts_sorted = sorted(noun_counts[\"all\"].items(), key=lambda x: -x[1])\n",
    "verb_counts_sorted = sorted(verb_counts[\"all\"].items(), key=lambda x: -x[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "402e61fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "if HAS_SPACY:\n",
    "    print(f\"\"\"\n",
    "      # Annotations = {len(all_anns)}\n",
    "      # Takes Annotated = {num_takes_covered}\n",
    "      Unique Annotations Per Take = {dict(coverage_by_task[\"all\"])}\n",
    "      # Descriptions = {len(anns)}\n",
    "      Avg Narrations per Annotation = {desc_per_ann.mean():.3f} (std dev = {desc_per_ann.std():.3f})\n",
    "      # Sentences = {desc_stats_df.num_sents.sum()}\n",
    "      Avg Sentences per Description = {desc_stats_df.num_sents.mean():.3f} (std dev = {desc_stats_df.num_sents.std():.3f})\n",
    "      # Words = {desc_stats_df.num_words.sum()}\n",
    "      Avg Words per Sentence = {desc_stats_df.words_per_sentence.mean():.3f} (std dev = {desc_stats_df.words_per_sentence.std():.3f})\n",
    "      # Unique Nouns = {len(noun_counts_sorted)}\n",
    "      # Unique Verbs = {len(verb_counts_sorted)}\n",
    "      \"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a1b31e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = {\n",
    "    \"Category\": [],\n",
    "    \"Takes >= 1x Coverage\": [],\n",
    "    \"Takes >= 2x Coverage\": [],\n",
    "    \"Number of Descriptions\": [],\n",
    "    \"Descriptions Per Minute\": [],\n",
    "    \"Unique Nouns\": [],\n",
    "    \"Unique Verbs\": [],\n",
    "}\n",
    "for cat in coverage_by_task.keys():\n",
    "    cat_name = cat\n",
    "    dpm = np.array(desc_per_minute[cat])\n",
    "    stats[\"Category\"].append(cat_name)\n",
    "    stats[\"Takes >= 1x Coverage\"].append(coverage_by_task[cat][1])\n",
    "    stats[\"Takes >= 2x Coverage\"].append(coverage_by_task[cat][2])\n",
    "    stats[\"Descriptions Per Minute\"].append(f\"{dpm.mean():.3f} (+- {dpm.std():.3f})\")\n",
    "    stats[\"Number of Descriptions\"].append(len(desc_by_task[cat]))\n",
    "    stats[\"Unique Nouns\"].append(len(noun_counts[cat]))\n",
    "    stats[\"Unique Verbs\"].append(len(verb_counts[cat]))\n",
    "\n",
    "stats_df = pd.DataFrame(stats)\n",
    "stats_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3f0d2c",
   "metadata": {},
   "source": [
    "## Embed the Descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dc4e4b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "embedder = SentenceTransformer('all-mpnet-base-v2', device=\"cuda\")\n",
    "\n",
    "def txt_simm(txt1, txt2):\n",
    "    query_embedding = embedder.encode(txt1)\n",
    "    passage_embedding = embedder.encode([txt2])\n",
    "    \n",
    "    return util.dot_score(query_embedding, passage_embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "989fcc4f-031f-4383-9452-f4bb87e6ce68",
   "metadata": {},
   "outputs": [],
   "source": [
    "embs_by_task = {}\n",
    "\n",
    "# NOTE:\n",
    "# this is a little redundant as \"all\" contains all parts\n",
    "# but this is fast enough to not matter\n",
    "for task, xs in tqdm(desc_by_task.items(), total=len(desc_by_task)):\n",
    "    txts = [x[\"text\"] for x in xs]\n",
    "    embs_by_task[task] = embedder.encode(txts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a039ebe3",
   "metadata": {},
   "source": [
    "## Cluster the Descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2c10140-3195-4949-aa19-1831163808ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import AgglomerativeClustering, KMeans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b879f2c-56a6-40d4-a35c-2dcfdf74ef6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cluster_task(task, num_clusters_or_threshold, cluster_type=\"kmeans\"):\n",
    "    embs = embs_by_task[task]\n",
    "    descs = desc_by_task[task]\n",
    "\n",
    "    # https://github.com/UKPLab/sentence-transformers/blob/master/examples/applications/clustering/kmeans.py\n",
    "    # num_clusters = 20\n",
    "    if cluster_type == \"kmeans\":\n",
    "        clustering_model = KMeans(n_clusters=num_clusters_or_threshold)\n",
    "    else:\n",
    "        clustering_model = AgglomerativeClustering(n_clusters=None, distance_threshold=num_clusters_or_threshold)\n",
    "        \n",
    "    clustering_model.fit(embs)\n",
    "    \n",
    "    cluster_assignment = clustering_model.labels_\n",
    "    # if isinstance(cluster_assignment, list):\n",
    "    #     cluster_assignment = {i: x for i, x in enumerate(cluster_assignment)}\n",
    "    \n",
    "    clustered_sentences = {}\n",
    "    for sentence_id, cluster_id in enumerate(cluster_assignment):\n",
    "        if cluster_id not in clustered_sentences:\n",
    "            clustered_sentences[cluster_id] = []\n",
    "        clustered_sentences[cluster_id].append(descs[sentence_id][\"text\"])\n",
    "    \n",
    "    for cluster_id, cluster in sorted(clustered_sentences.items(), key=lambda x: x[0]):\n",
    "        print(\"Cluster \", cluster_id)\n",
    "        print(random.sample(cluster, min(len(cluster), 5)), len(cluster))\n",
    "        print(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c5b911-f6c2-4d50-823c-1ebbf3c00dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_task(\"Soccer\", 20)\n",
    "# cluster_task(\"Soccer\", 1.5, \"aggl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b06406-6a07-4918-a355-9bc0655849f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cluster_task(\"Cooking\", 20)\n",
    "# cluster_task(\"Cooking\", 1.5, \"aggl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ebd9a2-b8d6-43d2-8923-7f93842706c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cluster_task(\"Bike Repair\", 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26301969-e84e-4b90-bda8-1b0d5b855d19",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# cluster_task(\"Health\", 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a00b73f9-ebdb-4d7f-afaa-c677c75f38b9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# cluster_task(\"Basketball\", 1.25, \"aggl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a395cfea",
   "metadata": {},
   "source": [
    "## Embedding Visualization & Dimension Reduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6043a324-e3ee-4539-9049-cf77fb00a027",
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.express as px\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "from PIL import Image\n",
    "from sklearn import preprocessing\n",
    "\n",
    "USING_SKLEARN = True\n",
    "PREFER_SKLEARN = True\n",
    "\n",
    "if PREFER_SKLEARN:\n",
    "    from sklearn.manifold import TSNE\n",
    "    print(\"Using sklearn\")\n",
    "else:\n",
    "    try:\n",
    "        # https://pypi.org/project/tsne-torch/\n",
    "        # https://github.com/CannyLab/tsne-cuda/blob/master/INSTALL.md\n",
    "        from tsnecuda import TSNE\n",
    "        USING_SKLEARN = False\n",
    "        print(\"Using CannyLab's tsnecuda\")\n",
    "    except Execption:\n",
    "        from sklearn.manifold import TSNE\n",
    "        print(f\"WARN: sklearn TSNE is used and not preferred: {traceback.format_exc()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bf6ef41",
   "metadata": {},
   "outputs": [],
   "source": [
    "task = \"Soccer\"  # NOTE: changeme to what you want to cluster\n",
    "embs = embs_by_task[task]\n",
    "descs = desc_by_task[task]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "323ea548-ab1f-4463-97f8-b96193ca8dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = time.time()\n",
    "X = embs_by_task[task]\n",
    "X_norm = preprocessing.MinMaxScaler(feature_range=(-1, 1)).fit_transform(X)\n",
    "X_tsne = TSNE(\n",
    "    n_components=2,\n",
    "    verbose=1,\n",
    "    n_iter=5000,\n",
    "    perplexity=300.0, # NOTE: adjust me\n",
    ").fit_transform(X_norm)\n",
    "t2 = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f69dd5f-e01c-4f00-a140-c5b77b51fce4",
   "metadata": {},
   "outputs": [],
   "source": [
    "t2 - t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d318660-274d-4ca5-bfc8-855b1edc3198",
   "metadata": {},
   "outputs": [],
   "source": [
    "xys = X_tsne.tolist()\n",
    "data_df = {\n",
    "    \"x\": [x for x, _ in xys],\n",
    "    \"y\": [y for _, y in xys],\n",
    "    \"description\": [x[\"text\"] for x in descs],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c32075b-054c-413c-931c-dc96388a9a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "px.scatter(data_df, x=\"x\", y=\"y\", color=None, hover_data=[\"description\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ca0685",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "expert_commentary",
   "language": "python",
   "name": "expert_commentary"
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}