notebooks/plot_Ego4D_user_split_stats.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b994c16",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Splits Ego4D in train/test/pretrain user-splits.\n",
    "Then, plots are generated of the video length in minutes (y-axis) per user (x-axis).\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "import datetime\n",
    "import json\n",
    "import os\n",
    "import os.path as osp\n",
    "import sys\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "plt.rcParams['font.size'] = 16\n",
    "# plt.rcParams.update(**font)\n",
    "\n",
    "# Use latex in mpl\n",
    "plt.rcParams['text.usetex'] = True\n",
    "plt.rcParams['text.latex.preamble'] = r'\\usepackage{amsmath}' #for \\text command\n",
    "\n",
    "parser = argparse.ArgumentParser(description=\"User split for ego4d LTA task.\")\n",
    "parser.add_argument(\n",
    "    \"--nb_users_thresh\",\n",
    "    help=\"Number users to keep as subset\",\n",
    "    default=50,\n",
    "    type=int,\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--nb_users_train\",\n",
    "    help=\"Number users to use for training, should be <= nb_users_thresh.\"\n",
    "         \"The 'nb_users_thresh - nb_users_train' are used for testing.\",\n",
    "    default=10,\n",
    "    type=int,\n",
    ")\n",
    "\n",
    "parser.add_argument(\n",
    "    \"--usersplit_criterion\",\n",
    "    help=\"Criterion to split the users on\",\n",
    "    default='random',\n",
    "    choices=['random', 'time_weighed'],\n",
    "    type=str,\n",
    ")\n",
    "\n",
    "parser.add_argument(\n",
    "    \"--user_videotime_min_thresh\",\n",
    "    help=\"Number of videominutes users need to remain in the subset\",\n",
    "    default=None,\n",
    "    type=int,\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--sort_by_col\",\n",
    "    help=\"Column in dataframe to sort on. (Default: total sum of user clip-video length)\",\n",
    "    default=\"sum_clip_duration_min\",\n",
    "    type=str,\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--p_output_dir\",\n",
    "    help=\"Parent dir to output timestamped dir including plots and json splits.\",\n",
    "    default=\"../imgs/notebooks_new\",\n",
    "    type=str,\n",
    ")\n",
    "parser.add_argument(\n",
    "    \"--seed\",\n",
    "    help=\"Seed numpy for deterministic splits\",\n",
    "    default=0,\n",
    "    type=int,\n",
    ")\n",
    "\n",
    "\n",
    "# Write both to file and stdout\n",
    "class Logger(object):\n",
    "    def __init__(self, filename=\"Default.log\"):\n",
    "        self.terminal = sys.stdout\n",
    "        self.log = open(filename, \"a\")\n",
    "\n",
    "    def write(self, message):\n",
    "        self.terminal.write(message)\n",
    "        self.log.write(message)\n",
    "\n",
    "    def flush(self):\n",
    "        self.log.flush()\n",
    "\n",
    "\n",
    "def generate_usersplit_from_trainval(user_id_col=\"fb_participant_id\"):\n",
    "    \"\"\" Get mini-train, test, and pretrain data splits.\"\"\"\n",
    "    args = parser.parse_args(args=[])\n",
    "\n",
    "    nb_users_test = args.nb_users_thresh - args.nb_users_train\n",
    "    train_ratio = args.nb_users_train / args.nb_users_thresh\n",
    "    np.random.seed(args.seed)\n",
    "\n",
    "    # META DATA\n",
    "    meta_data_file_path = \"../data/Ego4D/ego4d.json\"\n",
    "\n",
    "    # ANNOTATION DATA LTA\n",
    "    annotation_file_dir = \"../data/Ego4D/v1/annotations\"\n",
    "    annotation_file_names = {'train': \"fho_lta_train.json\", 'val': 'fho_lta_val.json',\n",
    "                             'test': 'fho_lta_test_unannotated.json'}\n",
    "    train_annotation_file = osp.join(annotation_file_dir, annotation_file_names['train'])\n",
    "    val_annotation_file = osp.join(annotation_file_dir, annotation_file_names['val'])\n",
    "\n",
    "    # Paths and logger\n",
    "    # Check outdir path\n",
    "    now = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')\n",
    "    output_dir = osp.join(args.p_output_dir, f\"{now}_ego4d_LTA_usersplit\")\n",
    "    os.makedirs(output_dir, exist_ok=True, mode=0o777)\n",
    "    sys.stdout = Logger(osp.join(output_dir, \"logger_dump.txt\"))\n",
    "\n",
    "    # check args\n",
    "    assert args.nb_users_thresh is None or args.user_videotime_min_thresh is None, \\\n",
    "        \"Can only define one thresholding method!\"\n",
    "\n",
    "    # Open meta data object\n",
    "    with open(meta_data_file_path, 'r') as meta_data_file:\n",
    "        meta_data_obj = json.load(meta_data_file)\n",
    "    meta_df = pd.json_normalize(meta_data_obj['videos'])\n",
    "    print(f\"meta_data.shape={meta_df.shape}\")\n",
    "\n",
    "    # Open train and val objects\n",
    "    with open(train_annotation_file, 'r') as train_file, \\\n",
    "            open(val_annotation_file, 'r') as val_file:\n",
    "        train_clips = json.load(train_file)['clips']\n",
    "        val_clips = json.load(val_file)['clips']\n",
    "\n",
    "    train_clips_df = pd.json_normalize(train_clips)\n",
    "    val_clips_df = pd.json_normalize(val_clips)\n",
    "    print(f\"trainshape={train_clips_df.shape}, valshape={val_clips_df.shape}\")\n",
    "\n",
    "    # Show overlapping\n",
    "    print(f\"Meta colnames={list(meta_df)}\")\n",
    "    print(f\"Annotation colnames={list(train_clips_df)}\")\n",
    "    overlapping_colnames = [x for x in list(meta_df) if x in list(train_clips_df)]\n",
    "    print(f\"Overlapping colnames={overlapping_colnames}\")\n",
    "\n",
    "    # MERGE dataframes on video_uid (Right join: keep annotation entries, but add video_uid info)\n",
    "    train_joined_df = pd.merge(meta_df, train_clips_df,\n",
    "                               on=\"video_uid\", validate=\"one_to_many\", how=\"right\")\n",
    "    val_joined_df = pd.merge(meta_df, val_clips_df,\n",
    "                             on=\"video_uid\", validate=\"one_to_many\", how=\"right\")\n",
    "    print(f\"train_joined_df={train_joined_df.shape}, val_joined_df={val_joined_df.shape}\")\n",
    "\n",
    "    # CONCAT the dataframes (312 rows × 12 columns)\n",
    "    trainval_joined_df = pd.concat([train_joined_df, val_joined_df], ignore_index=True, sort=False)\n",
    "\n",
    "    # FIND USERS that satisfy video-length threshold\n",
    "    # Note: video_uid relates to the entire uncut raw video,\n",
    "    # these are split into ~5-min clips, denoted with clip_id for the annotations.\n",
    "    trainval_user_df, nan_user_df = summarize_clips_by_user(trainval_joined_df)\n",
    "\n",
    "    # Sort users on sum_length\n",
    "    trainval_user_df = trainval_user_df.sort_values(by=[args.sort_by_col], ascending=False)\n",
    "\n",
    "    pretrain_user_ids = []\n",
    "    pretrain_sort_values = []\n",
    "    # Keep only highest in sorted\n",
    "    if args.nb_users_thresh is not None:\n",
    "        print(f\"Thresholding on nb_users_thresh\")\n",
    "        sorted_user_ids = trainval_user_df[: args.nb_users_thresh][user_id_col].tolist()\n",
    "        user_sort_values = trainval_user_df[: args.nb_users_thresh][args.sort_by_col].tolist()\n",
    "\n",
    "        # PRETRAIN: Leftover users + NaN user\n",
    "        pretrain_user_ids = trainval_user_df[args.nb_users_thresh:][user_id_col].tolist()\n",
    "        pretrain_sort_values = trainval_user_df[args.nb_users_thresh:][args.sort_by_col].tolist()\n",
    "\n",
    "        # Add NaN user\n",
    "        pretrain_user_ids.append(None)\n",
    "        pretrain_sort_values.extend(nan_user_df[args.sort_by_col].tolist())\n",
    "\n",
    "    elif args.user_videotime_min_thresh is not None:\n",
    "        raise NotImplementedError(\"Not supporting this split\")\n",
    "        print(f\"Thresholding on user_videotime_min_thresh\")\n",
    "        user_subset_df = trainval_user_df.loc[trainval_user_df[args.sort_by_col] >= args.user_videotime_min_thresh]\n",
    "        sorted_user_ids = user_subset_df[user_id_col].tolist()\n",
    "        user_sort_values = user_subset_df[args.sort_by_col].tolist()\n",
    "\n",
    "    else:  # No cutoff\n",
    "        print(f\"No thresholding = no pretrain data\")\n",
    "        sorted_user_ids = trainval_user_df[user_id_col].tolist()\n",
    "        user_sort_values = trainval_user_df[args.sort_by_col].tolist()\n",
    "\n",
    "    # Get train/test splits from train/val datasets\n",
    "    if args.usersplit_criterion == 'random':\n",
    "        shuffled_idxs = np.random.permutation(np.arange(len(sorted_user_ids)))\n",
    "        train_idxs, test_idxs = shuffled_idxs[:args.nb_users_train], shuffled_idxs[args.nb_users_train:]\n",
    "\n",
    "    elif args.usersplit_criterion == 'time_weighed':\n",
    "        raise NotImplementedError()\n",
    "    else:\n",
    "        raise ValueError()\n",
    "\n",
    "    sorted_user_ids, user_sort_values = np.array(sorted_user_ids), np.array(user_sort_values)\n",
    "    train_user_ids, test_user_ids = sorted_user_ids[train_idxs], sorted_user_ids[test_idxs]\n",
    "    train_sort_values, test_sort_values = user_sort_values[train_idxs], user_sort_values[test_idxs]\n",
    "\n",
    "    # Print summary\n",
    "    nb_traintest_users_subset = len(sorted_user_ids)\n",
    "    nb_train_users_subset = len(train_user_ids)\n",
    "    nb_test_users_subset = len(test_user_ids)\n",
    "    nb_pretrain_users_subset = len(pretrain_user_ids)\n",
    "    nb_total_users = nb_pretrain_users_subset + nb_traintest_users_subset\n",
    "\n",
    "    print_summary(user_sort_values, nb_traintest_users_subset, nb_total_users, \"Train+Test user subset\")\n",
    "    print_summary(train_sort_values, nb_train_users_subset, nb_traintest_users_subset, \"Train user subset\")\n",
    "    print_summary(test_sort_values, nb_test_users_subset, nb_traintest_users_subset, \"Test user subset\")\n",
    "    print_summary(pretrain_sort_values, nb_pretrain_users_subset, nb_total_users, \"Pretrain user subset\")\n",
    "\n",
    "    # Summary plot (pdf?)\n",
    "    ylabel = 'Video length (minutes)'\n",
    "    xlabel = \"User\"\n",
    "    title = \"Sum of clip video lengths (min) per user\"\n",
    "\n",
    "    # Get colors for train+test summary plot\n",
    "    labels = [r'$\\mathcal{U}_\\text{train}$', r'$\\mathcal{U}_\\text{test}$']\n",
    "    include_nan_pretrain = False\n",
    "    bar_colors =[ sns.color_palette(\"rocket\")[1], sns.color_palette(\"rocket\")[5]]\n",
    "    y_axis = [list(train_sort_values), list(test_sort_values)]\n",
    "    x_axis = [train_idxs, test_idxs]\n",
    "    if len(pretrain_sort_values) > 0:\n",
    "        x_offset = nb_traintest_users_subset\n",
    "\n",
    "        if include_nan_pretrain:\n",
    "            y_axis.append(list(pretrain_sort_values))\n",
    "            x_axis.append([idx for idx in range(x_offset, x_offset + len(pretrain_sort_values))])\n",
    "        else:\n",
    "            y_axis.append(list(pretrain_sort_values[:-1]))\n",
    "            x_axis.append([idx for idx in range(x_offset, x_offset + len(pretrain_sort_values)-1)])\n",
    "        bar_colors.append(sns.color_palette(\"rocket\")[3])\n",
    "        labels.append(r'$\\mathcal{U}_\\text{population}$')\n",
    "    plot_tag = 'TRAIN_TEST_PRETRAIN'\n",
    "    plot_barchart(x_axis, y_axis, title=f'{plot_tag} {title}', ylabel=ylabel, legend_labels=labels,\n",
    "                  xlabel=xlabel,\n",
    "                  bar_colors=bar_colors, output_file=osp.join(output_dir, f'{plot_tag}_video_freq_plot.pdf'))\n",
    "\n",
    "#     y_axis = sorted(train_sort_values, reverse=True)\n",
    "#     x_axis = [idx for idx in range(len(train_sort_values))]\n",
    "#     plot_tag = 'TRAIN'\n",
    "#     plot_barchart([x_axis], [y_axis], title=f'{plot_tag} {title}', ylabel=ylabel, xlabel=xlabel,\n",
    "#                   output_file=osp.join(output_dir, f'{plot_tag}_video_freq_plot.pdf'))\n",
    "\n",
    "#     y_axis = sorted(test_sort_values, reverse=True)\n",
    "#     x_axis = [idx for idx in range(len(test_sort_values))]\n",
    "#     plot_tag = 'TEST'\n",
    "#     plot_barchart([x_axis], [y_axis], title=f'{plot_tag} {title}', ylabel=ylabel, xlabel=xlabel,\n",
    "#                   output_file=osp.join(output_dir, f'{plot_tag}_video_freq_plot.pdf'))\n",
    "\n",
    "#     y_axis = sorted(pretrain_sort_values, reverse=True)\n",
    "#     x_axis = [idx for idx in range(len(pretrain_sort_values))]\n",
    "#     plot_tag = 'PRETRAIN'\n",
    "#     plot_barchart([x_axis], [y_axis], title=f'{plot_tag} {title}', ylabel=ylabel, xlabel=xlabel,\n",
    "#                   output_file=osp.join(output_dir, f'{plot_tag}_video_freq_plot.pdf'))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def print_summary(user_sort_values, nb_users_subset, nb_total_users, title: str):\n",
    "    print(f\"\\n{'*' * 20} SUMMARY: {title} {'*' * 20}\")\n",
    "    print(f\"Retaining a total of {nb_users_subset} / {nb_total_users} users\")\n",
    "    print(f\"SUM = {sum(user_sort_values)}\")\n",
    "    print(f\"MAX = {max(user_sort_values)}, MIN = {min(user_sort_values)}\")\n",
    "    print(f\"HEAD = {user_sort_values[:10]}...\")\n",
    "    print(f\"TAIL = ...{user_sort_values[-10:]}\")\n",
    "    print(f\"{'*' * 50}\")\n",
    "\n",
    "\n",
    "def summarize_clips_by_user(joined_df):\n",
    "    \"\"\"Group annotation entries by clip_uid, then group those unique clips by user.\"\"\"\n",
    "    clip_df = joined_df.groupby(joined_df['clip_uid'], as_index=False).agg(\n",
    "        {'fb_participant_id': lambda x: np.unique(x).tolist(),\n",
    "         'scenarios': list,\n",
    "         'verb': list, 'noun': list, 'verb_label': list, 'noun_label': list, 'action_idx': list,\n",
    "         #          'video_uid':list,'duration_sec':list, # This is the raw uncut video, don't need this info\n",
    "         'clip_id': list, 'clip_parent_start_sec': lambda x: np.unique(x).tolist(),\n",
    "         'clip_parent_end_sec': lambda x: np.unique(x).tolist()})\n",
    "\n",
    "    # Check users/clip_starts and ends are only 1 unique\n",
    "    assert (clip_df.fb_participant_id.apply(len) == 1).all()\n",
    "    assert (clip_df.clip_parent_start_sec.apply(len) == 1).all()\n",
    "    assert (clip_df.clip_parent_end_sec.apply(len) == 1).all()\n",
    "\n",
    "    # Unpack\n",
    "    for col_name in ['fb_participant_id', 'clip_parent_start_sec', 'clip_parent_end_sec']:\n",
    "        clip_df[col_name] = clip_df[col_name].apply(lambda x: x[0])\n",
    "\n",
    "    # Get actual clip lengths in seconds (~5min=300s)\n",
    "    clip_df['clip_duration_sec'] = clip_df.loc[:, ('clip_parent_end_sec', 'clip_parent_start_sec')].apply(\n",
    "        lambda x: x[0] - x[1], axis=1)\n",
    "\n",
    "    # Group by fb_participant_id, which has allocated multiple 5min clips (unique clip_uid's)\n",
    "    user_df = clip_df.groupby(clip_df['fb_participant_id'], as_index=False, dropna=False).agg(  # NaN Dropped by default\n",
    "        {\n",
    "            'scenarios': list,\n",
    "            'verb': list, 'noun': list, 'verb_label': list, 'noun_label': list, 'action_idx': list,\n",
    "            'clip_id': list, 'clip_parent_start_sec': list, 'clip_parent_end_sec': list, 'clip_duration_sec': list}\n",
    "    )\n",
    "\n",
    "    # Sum clip lengths per user\n",
    "    user_df['sum_clip_duration_sec'] = user_df['clip_duration_sec'].apply(sum)\n",
    "    user_df['sum_clip_duration_min'] = user_df['sum_clip_duration_sec'].apply(lambda x: x / 60)\n",
    "\n",
    "    # The scenarios only apply to the raw uncut video, not the 5min clips\n",
    "    user_df = user_df.rename(columns={\"scenarios\": \"possible_clip_scenarios\"})\n",
    "\n",
    "    # Extract NaN user entry\n",
    "    nan_user_df = user_df[user_df.fb_participant_id.isnull()]\n",
    "    user_df = user_df[user_df.fb_participant_id.notnull()]\n",
    "\n",
    "    # Check that no NaN user for valid users\n",
    "    assert not (user_df['fb_participant_id'].isna().any())\n",
    "\n",
    "    return user_df, nan_user_df\n",
    "\n",
    "\n",
    "def plot_barchart(x_axis: list[list], y_vals: list[list], title, ylabel, xlabel, y_labels=None, x_labels=None,\n",
    "                  legend_labels=None,\n",
    "                  grid=False, yerror: list[list] = None, bar_align='edge', bar_colors=None,\n",
    "                  figsize=(8, 3), log=False, interactive=False, x_minor_ticks=None, output_file=None):\n",
    "    max_val = max(el for y_l in y_vals for el in y_l)\n",
    "    my_cmap = plt.get_cmap(\"plasma\")\n",
    "    fig = plt.figure(figsize=figsize, dpi=600)  # So all bars are visible!\n",
    "    ax = plt.subplot()\n",
    "\n",
    "    for plot_idx, (x, y) in enumerate(zip(x_axis, y_vals)):\n",
    "        plot_yerror = None if yerror is None else yerror[plot_idx]\n",
    "        bar_color = my_cmap.colors if bar_colors is None else bar_colors[plot_idx]\n",
    "        legend_label = None if legend_labels is None else legend_labels[plot_idx]\n",
    "        plt.bar(x, height=y, yerr=plot_yerror, label=legend_label,\n",
    "                color=bar_color, align=bar_align, width=0.7, log=log,\n",
    "                alpha=0.8\n",
    "#                linewidth=0.1,edgecolor='white'\n",
    "               )\n",
    "\n",
    "        x_center = [el+0.4 for el in x]\n",
    "        plt.plot(x_center, y, marker=\"v\",markersize=0.8,\n",
    "                 linestyle=\"\", alpha=1, color=\"black\")\n",
    "\n",
    "    if legend_labels is not None:\n",
    "        plt.legend(prop={'size': 18})\n",
    "\n",
    "\n",
    "    if x_minor_ticks is not None:\n",
    "        ax.set_xticks(x_minor_ticks, minor=True)\n",
    "\n",
    "    # Hide the right and top spines\n",
    "    ax.spines.right.set_visible(False)\n",
    "    ax.spines.top.set_visible(False)\n",
    "\n",
    "    # Only show ticks on the left and bottom spines\n",
    "    ax.yaxis.set_ticks_position('left')\n",
    "    ax.xaxis.set_ticks_position('bottom')\n",
    "\n",
    "    if x_labels:\n",
    "        plt.xticks(x_axis, x_labels, rotation='vertical')\n",
    "    if y_labels:\n",
    "        plt.yticks(y_vals, y_labels)\n",
    "\n",
    "    plt.ylim(None, max_val * 1.01)\n",
    "    plt.xlim(-4, 199)\n",
    "    plt.xlabel(xlabel)\n",
    "    plt.ylabel(ylabel)\n",
    "    print(title)\n",
    "#     plt.title(title)\n",
    "    plt.grid(grid, which='both')\n",
    "\n",
    "    # Save\n",
    "    fig.tight_layout()\n",
    "    if output_file is not None:\n",
    "        fig.savefig(output_file, bbox_inches='tight')\n",
    "        print(f\"Saved plot: {output_file}\")\n",
    "\n",
    "    plt.show()\n",
    "    plt.clf()\n",
    "\n",
    "\n",
    "generate_usersplit_from_trainval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b597ef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}