04_analyze.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# analyze\n",
    "\n",
    "> Tools for computing movement/pose similarity time series, clustering\n",
    "  for single- and multi-dancer videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp analyze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import openpifpaf\n",
    "from openpifpaf.plugins.coco.constants import COCO_KEYPOINTS, COCO_PERSON_SKELETON\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from skbio.stats.distance import mantel\n",
    "from scipy.spatial.distance import squareform\n",
    "from scipy.interpolate import interp1d\n",
    "from sklearn.cluster import OPTICS\n",
    "from Bio import pairwise2\n",
    "\n",
    "from choreo_k.matrixify import get_pose_matrix, get_normalized_coords, compare_poses_cosine, get_laplacian_matrix\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\n",
    "  action='ignore', module='matplotlib.figure', category=UserWarning,\n",
    "  message=('This figure includes Axes that are not compatible with tight_layout, '\n",
    "           'so results might be incorrect.'))\n",
    "\n",
    "def smooth_series(x, window_len=11, window='flat'):\n",
    "    \"\"\" Smooth a time series via a sliding window average\n",
    "        From https://scipy-cookbook.readthedocs.io/items/SignalSmooth.html\n",
    "    \"\"\"\n",
    "    if x.ndim != 1:\n",
    "        raise(ValueError, \"smooth only accepts 1 dimension arrays.\")\n",
    "\n",
    "    if x.size < window_len:\n",
    "        raise(ValueError, \"Input vector needs to be bigger than window size.\")\n",
    "\n",
    "    if window_len<3:\n",
    "        print(\"WARNING: window length too small for smoothing, returning input data\")\n",
    "        return x\n",
    "\n",
    "    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:\n",
    "        raise(ValueError, \"Window is one of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'\")\n",
    "\n",
    "    s=np.r_[x[window_len-1:0:-1],x,x[-2:-window_len-1:-1]]\n",
    "    if window == 'flat': #moving average\n",
    "        w=np.ones(window_len,'d')\n",
    "    else:\n",
    "        w=eval('np.'+window+'(window_len)')\n",
    "\n",
    "    y=np.convolve(w/w.sum(),s,mode='valid')\n",
    "\n",
    "    # Move the window 1/2 width back to avoid lag\n",
    "    if window_len % 2 == 0:\n",
    "        return y[int(window_len/2)-1:-(int(window_len/2))]\n",
    "    else:\n",
    "        return y[int(window_len/2):-(int(window_len/2))]\n",
    "\n",
    "\n",
    "def corr_time_series_matrix(pose_data, method='distance'):\n",
    "    \"\"\" Generate a full time-series pose similarity heatmap for all available\n",
    "        poses and frames from the video. This code can use either pose\n",
    "        characterization approach; in practice, the distance matrix-based analyses\n",
    "        take longer to calculate but are more accurate.\n",
    "    \"\"\"\n",
    "    pose_correlations = []\n",
    "    for i, pi in enumerate(pose_data):\n",
    "        print(\"Comparing frame\",i,\"to the rest\")\n",
    "        corr_row = []\n",
    "        if method == 'distance':\n",
    "            mi = get_pose_matrix(pi)\n",
    "        else: # method == 'laplacian'\n",
    "            mi = get_laplacian_matrix(pi, figure_type='flipped_figures')\n",
    "        for j, pj in enumerate(pose_data):\n",
    "            if j < i:\n",
    "                corr_row.append(pose_correlations[j][i])\n",
    "            elif j == i:\n",
    "                corr_row.append(float(1))\n",
    "            else:\n",
    "                if mi is None:\n",
    "                    corr_row.append(float(0))\n",
    "                elif method == 'distance':\n",
    "                    mj = get_pose_matrix(pj)\n",
    "                    if mj is None:\n",
    "                        corr_row.append(float(0))\n",
    "                    else:\n",
    "                        corr_row.append(mantel(mi, mj)[0])\n",
    "                else: # method == 'laplacian'\n",
    "                    mj = get_laplacian_matrix(pj, figure_index=0, figure_type='flipped_figures')\n",
    "                    if mj is None:\n",
    "                        corr_row.append(float(0))\n",
    "                    else:\n",
    "                        corr_row.append(1 - abs(np.subtract(mi.todense(), mj.todense()).sum()))\n",
    "        pose_correlations.append(corr_row)\n",
    "\n",
    "    return pose_correlations\n",
    "\n",
    "\n",
    "def correlate_time_series(pose_data1, pose_data2, method='correlate', figure_type='figures'):\n",
    "    if pose_data2 is None:\n",
    "        pose_data2 = pose_data1\n",
    "    pose_correlations = []\n",
    "    for i, pi in enumerate(pose_data1):\n",
    "        print(\"Comparing frame\",i,\"to the rest\")\n",
    "        corr_row = []\n",
    "        if method == 'correlate': # Distance matrix correlation\n",
    "            mi = get_pose_matrix(pi, figure_type=figure_type)\n",
    "        elif method == 'distance': # Cosine similarity of normalized keypoints\n",
    "            mi = get_normalized_coords(pi, figure_type=figure_type)\n",
    "        else:\n",
    "            mi = get_laplacian_matrix(pi, figure_type=figure_type)\n",
    "        for j, pj in enumerate(pose_data2):\n",
    "            if mi is None:\n",
    "                corr_row.append(float(0))\n",
    "            elif method == 'correlate':\n",
    "                mj = get_pose_matrix(pj, figure_type=figure_type)\n",
    "                if mj is None:\n",
    "                    corr_row.append(float(0))\n",
    "                else:\n",
    "                    corr_row.append(mantel(mi, mj)[0])\n",
    "            elif method == 'distance':\n",
    "                mj = get_normalized_coords(pj, figure_type=figure_type)\n",
    "                if mj is None:\n",
    "                    corr_row.append(float(0))\n",
    "                else:\n",
    "                    corr_row.append(compare_poses_cosine(mi, mj))\n",
    "            else:\n",
    "                mj = get_laplacian_matrix(pj, figure_index=0, figure_type=figure_type)\n",
    "                if mj is None:\n",
    "                    corr_row.append(float(0))\n",
    "                else:\n",
    "                    corr_row.append(1 - abs(np.subtract(mi.todense(), mj.todense()).sum()))\n",
    "        pose_correlations.append(corr_row)\n",
    "        \n",
    "    return pose_correlations\n",
    "\n",
    "\n",
    "def fill_nans_scipy1(padata, pkind='linear'):\n",
    "    \"\"\" Fill in missing values from a time series, after the first non-NAN value\n",
    "        and up to the last non-NAN value. Note that scipy.interpolated.interp1d\n",
    "        provides a lot more options (splines, quadratic, etc.)\n",
    "    \"\"\"\n",
    "    \n",
    "    aindexes = np.arange(padata.shape[0])\n",
    "    agood_indexes, = np.where(np.isfinite(padata))\n",
    "\n",
    "    first_non_nan_index = agood_indexes[0] \n",
    "    last_non_nan_index = agood_indexes[-1]\n",
    "\n",
    "    f = interp1d(agood_indexes\n",
    "           , padata[agood_indexes]\n",
    "           , bounds_error=False\n",
    "           , copy=False\n",
    "           , fill_value=\"extrapolate\"\n",
    "           , kind=pkind)\n",
    "\n",
    "    interpolated = f(aindexes)\n",
    "    interpolated[0:first_non_nan_index] = np.nan\n",
    "    interpolated[last_non_nan_index+1:] = np.nan\n",
    "    return interpolated\n",
    "\n",
    "\n",
    "def movements_time_series(pose_data, pose_index=-1, method='distance', figure_type='flipped_figures', video_file=None):\n",
    "    \"\"\" Calculate a time series of the differences between each pair of poses in a\n",
    "        sequence. This works with a single figure (pose_index=0) or all the figures\n",
    "        in the video (pose_index=-1). It can be run on its own, but typically this is\n",
    "        a helper function for process_movement_series() (below).\n",
    "    \"\"\"\n",
    "    \n",
    "    per_frame_movements = []\n",
    "    frame_timecodes = []\n",
    "\n",
    "    pose_indices = []\n",
    "\n",
    "    max_figures, total_time, total_figures = count_figures_and_time(pose_data, figure_type)\n",
    "\n",
    "    threshold = .7\n",
    "\n",
    "    #print(\"FIGURES PER FRAME IN TIME SERIES:\",max_figures)\n",
    "\n",
    "    # Typically the pose index is only specified if you know there's only one dancer\n",
    "    # (in which case it's always 0)\n",
    "    if pose_index != -1:\n",
    "        max_figures = 1\n",
    "\n",
    "    for f, frame in enumerate(pose_data):\n",
    "        frame_movements = []\n",
    "        frame_timecodes.append(frame['time'])\n",
    "        if f < len(pose_data)-1:\n",
    "            for p in range(max_figures):\n",
    "                this_motion = np.nan\n",
    "                movers = np.array([])\n",
    "                # figure p must be available for both f and f-1\n",
    "                # NOTE check p < len(pose_data[f][figure_type]) in case the data for that frame has been\n",
    "                # truncated to [] (for example if extraneous data from the end of the video has been removed)\n",
    "                if p < len(pose_data[f-1][figure_type]) and p < len(pose_data[f][figure_type]) and pose_data[f-1][figure_type][p].data.shape[0] != 0 and pose_data[f][figure_type][p].data.shape[0] != 0:\n",
    "                    p1_conf = sum([c[2] for c in pose_data[f-1][figure_type][p].data]) / float(len(pose_data[f-1][figure_type][p].data))\n",
    "                    p2_conf = sum([c[2] for c in pose_data[f][figure_type][p].data]) / float(len(pose_data[f][figure_type][p].data))\n",
    "                    # XXX USE A BETTER CRITERION FOR SKIPPING POSES IF CONFIDENCE IS LOW\n",
    "                    if p1_conf > threshold and p2_conf > threshold:\n",
    "                        if method == 'distance':\n",
    "                            plot_type = 'distance'\n",
    "                            dmatrix1 = squareform(get_pose_matrix(pose_data[f-1], p, figure_type))\n",
    "                            dmatrix2 = squareform(get_pose_matrix(pose_data[f], p, figure_type))\n",
    "                            diffmatrix = np.absolute(dmatrix1 - dmatrix2)\n",
    "                            movers = diffmatrix.sum(axis=1)\n",
    "                            this_motion = movers.sum(axis=0) # For debugging\n",
    "                        else:\n",
    "                            plot_type = 'delaunay'\n",
    "                            # Per-keypoint movements are not useful for Laplacian comparisons\n",
    "                            similarity = compare_laplacians(pose_data[f-1], pose_data[f], p, figure_type)\n",
    "                            # Can we get meaningful movement values if laplacians are of different sizes?\n",
    "                            if similarity is not None:\n",
    "                                movers = np.array([1 - similarity])\n",
    "                \n",
    "                frame_movements.append(movers)    \n",
    "        \n",
    "        per_frame_movements.append(frame_movements)\n",
    "\n",
    "    return [per_frame_movements, frame_timecodes, max_figures]\n",
    "\n",
    "\n",
    "def process_movement_series(pose_data, pose_index=-1, figure_type='flipped_figures', video_file=None, method='distance', interpolate=True, viz=True):\n",
    "    \"\"\" Smooth, summarize, visualize movement data for one or more figures across a\n",
    "        time series.\n",
    "        Also visualize aggregate movement data for each keypoint, if distance matrix\n",
    "        method is used.\n",
    "    \"\"\"\n",
    "    \n",
    "    print(\"GETTING MOVEMENT TIME SERIES\")\n",
    "    per_frame_results, frame_timecodes, max_figures = movements_time_series(pose_data, pose_index, method, figure_type, video_file)\n",
    "\n",
    "    print(\"CALCULATING CHARACTERISTICS OF TIME SERIES\")\n",
    "\n",
    "    window_length = 5\n",
    "    if video_file is not None:\n",
    "        fps, total_frames = get_video_stats(video_file)\n",
    "        window_length = max(window_length, int(round(fps/2.0)))\n",
    "\n",
    "    movers = [] # To get the aggregate avg movement of each keypoint (not for Laplacian)\n",
    "    movement_series = []\n",
    "    frame_times = []\n",
    "\n",
    "    smoothed_movement_series = []\n",
    "\n",
    "    for j in range(max_figures):\n",
    "        movement_series.append([])  \n",
    "        smoothed_movement_series.append([])\n",
    "\n",
    "    per_frame_movements = []\n",
    "    \n",
    "    for f, frame in enumerate(per_frame_results):\n",
    "        frame_movements = np.zeros(TOTAL_COORDS)\n",
    "        frame_times.append(frame_timecodes[f])\n",
    "        for j in range(max_figures):\n",
    "            if j >= len(frame) or frame[j].shape[0] == 0:\n",
    "                #print(\"POSE\",j,\"HAS NO MOVEMENT DATA\")\n",
    "                movement_series[j].append(np.nan)\n",
    "            elif method == 'distance':\n",
    "                frame_movements = np.add(frame_movements, frame[j])\n",
    "                movement_series[j].append(sum(frame[j]))\n",
    "                movers.append(np.array(frame[j]))\n",
    "            else: # method == 'laplacian'\n",
    "                movement_series[j].append(frame[j][0])\n",
    "            movers.append(frame[j])\n",
    "        per_frame_movements.append(frame_movements)\n",
    "\n",
    "    figure_time_series = np.array(movers)\n",
    "\n",
    "    # Not sure if there's a meaningful way to aggregate the per-keypoint data\n",
    "    # for the graph Laplacian approach (e.g., to be able to quantify how much\n",
    "    # each keypoint moved during the video).\n",
    "    if method == 'distance':\n",
    "        movement_means = np.nanmean(figure_time_series, axis=0)\n",
    "        movement_stdevs = np.nanstd(figure_time_series, axis=0)\n",
    "\n",
    "    # Window length is half of fps (or ~5, whichever is larger)\n",
    "    for j in range(max_figures):\n",
    "        if interpolate:\n",
    "            if np.isnan(np.nanmax(movement_series[j])):\n",
    "                #print(\"Movement series is all nans, skipping\")\n",
    "                continue\n",
    "            smoothed_movement_series[j] = smooth_series(fill_nans_scipy1(np.asarray(movement_series[j]), pkind='linear'),window_length).tolist()\n",
    "        else:\n",
    "            smoothed_movement_series[j] = smooth_series(np.array(movement_series[j]),window_length).tolist()\n",
    "\n",
    "    if viz:\n",
    "        print(\"VISUALIZING TIME SERIES CHARACTERISTICS\")\n",
    "    \n",
    "        if method == 'distance':\n",
    "            plt.figure()\n",
    "            plt.xticks(np.arange(TOTAL_COORDS))\n",
    "\n",
    "            plt.bar(np.arange(TOTAL_COORDS), movement_means)# yerr=movement_stdevs)\n",
    "\n",
    "        fig = plt.figure(figsize=(12,6), constrained_layout=True)\n",
    "        fig.dpi=100\n",
    "\n",
    "        for j in range(len(smoothed_movement_series)):\n",
    "            if (len(smoothed_movement_series[j]) > 0) and not np.isnan(np.nanmax(smoothed_movement_series[j])):\n",
    "                plt.plot(frame_times,smoothed_movement_series[j]) #,color)\n",
    "        plt.show()\n",
    "\n",
    "    if method == 'distance':\n",
    "        return [smoothed_movement_series, frame_times, per_frame_movements, movement_means, movement_stdevs]\n",
    "    else:\n",
    "        return [smoothed_movement_series, frame_times]\n",
    "    \n",
    "\n",
    "def average_poses(pose_series, descriptors, source_figures='zeroified_figures', flip=True):\n",
    "    # Descriptors are [[frame_index, pose_index] ...]\n",
    "    # XXX Add an option to average all frames/poses across a range, or an entire video?\n",
    "    all_poses = []\n",
    "    for descriptor in descriptors:\n",
    "        all_poses.append(pose_series[descriptor[0]][source_figures][descriptor[1]].data)\n",
    "    print(len(all_poses))\n",
    "    poses_array = np.array(all_poses)\n",
    "    avg_array = np.sum(poses_array, axis=0)/len(poses_array)\n",
    "    this_annotation = openpifpaf.Annotation(keypoints=COCO_KEYPOINTS, skeleton=COCO_PERSON_SKELETON).set(avg_array, fixed_score=None)\n",
    "    #this_annotation = openpifpaf.decoder.annotation.Annotation(keypoints=COCO_KEYPOINTS, skeleton=COCO_PERSON_SKELETON)\n",
    "    for f, xyv in enumerate(avg_array):\n",
    "        this_annotation.add(f, xyv)\n",
    "    if flip:\n",
    "        this_annotation = flip_detections([this_annotation])[0]\n",
    "    return this_annotation\n",
    "\n",
    "def get_feature_vectors(pose_series, figure_type='aligned_figures', method='distance'):\n",
    "    # Convert matrices into feature vectors to send to the clustering algorithm\n",
    "    features = []\n",
    "    descriptors = []\n",
    "    for f, frame_info in enumerate(pose_series):\n",
    "        for p, pose_info in enumerate(frame_info[figure_type]):\n",
    "            if method == 'distance':\n",
    "                pose_matrix = get_pose_matrix(frame_info, p)\n",
    "            else: # method == 'laplacian'\n",
    "                pose_matrix = get_laplacian_matrix(frame_info, p)\n",
    "            if pose_matrix is not None:\n",
    "                features.append(pose_matrix)\n",
    "                descriptors.append([f,p])\n",
    "    return([features, descriptors])\n",
    "\n",
    "\n",
    "def cluster_poses(poses_series, figure_type='aligned_figures', min_samples=50):\n",
    "    # min_samples can be set according to some rule, e.g., a fraction or multiple of\n",
    "    # frames per second * median number of people in a frame\n",
    "    print(\"Getting feature vectors\")\n",
    "    [poses_features, descriptors] = get_feature_vectors(poses_series, figure_type)\n",
    "    data_array = np.array(poses_features)\n",
    "    print(data_array.shape)\n",
    "    print(len(descriptors))\n",
    "\n",
    "    print(\"Fitting OPTICS\")\n",
    "    #labels = DBSCAN(eps=100000).fit_predict(features_array)\n",
    "    labels = OPTICS(min_samples=min_samples, metric='sqeuclidean').fit_predict(data_array)\n",
    "\n",
    "    return [labels, descriptors]\n",
    "\n",
    "\n",
    "def get_cluster_averages_and_indices(labels, descriptors, pose_series, figure_type='figures', video_file=None, flip_figures=False):\n",
    "    \"\"\" Average the members of each pose cluster to get the representative \"average\"\n",
    "        pose for the cluster.\n",
    "    \"\"\"\n",
    "\n",
    "    label_keys = []\n",
    "    for label in labels:\n",
    "        if label != -1 and label not in label_keys:\n",
    "            label_keys.append(label)\n",
    "    label_keys.sort()\n",
    "\n",
    "    cluster_indices = {}\n",
    "    cluster_averages = {}\n",
    "    cluster_avg_poses = {}\n",
    "\n",
    "    total_poses = len(label_keys)\n",
    "\n",
    "    for label in label_keys:\n",
    "        indices = [j for j, x in enumerate(labels) if x == label]\n",
    "        descs = [descriptors[indices[k]] for k in range(len(indices))]\n",
    "        print(\"CLUSTER\",label,\"|\",len(indices),\"POSES\")\n",
    "        cluster_indices[label] = indices\n",
    "        print(descriptors[indices[0]],\"CLUSTER\",label,'FIRST POSE')\n",
    "        if video_file is not None:\n",
    "            fig = excerpt_pose(video_file, pose_series[descriptors[indices[0]][0]], descriptors[indices[0]][1], show=True, source_figure=figure_type, flip_figures=flip_figures)\n",
    "        avg_pose = average_poses(pose_series, descs)\n",
    "        cluster_averages[label] = matrixify_pose(avg_pose.data)\n",
    "        cluster_avg_poses[label] = avg_pose\n",
    "        print(\"CLUSTER\",label,'AVERAGE POSE')\n",
    "        plot_poses(avg_pose)\n",
    "\n",
    "    return [cluster_averages, cluster_indices, cluster_avg_poses]\n",
    "\n",
    "\n",
    "def find_nearest_pose(pose_matrix, cluster_averages):\n",
    "    best_corr = 0\n",
    "    best_label = -1\n",
    "    for label in cluster_averages:\n",
    "        corr = mantel(pose_matrix, cluster_averages[label])[0]\n",
    "        if corr > best_corr:\n",
    "            best_label = label\n",
    "            best_corr = corr\n",
    "    return best_label\n",
    "\n",
    "\n",
    "CELL_HEIGHT=120\n",
    "\n",
    "\n",
    "def render_pose_distribution(heatmap, poses_series, labels, descriptors, closest_matches=None, show=True, video_file=None, time_index=None, cell_height=CELL_HEIGHT, xlim=None):\n",
    "    \"\"\" Draw a pose cluster timeline based on a precomputed clustering and\n",
    "        assigment of non-clustered poses to clusters via compute_pose_distribution()\n",
    "        below. Passing in the heatmap from the previous step saves a lot of time.\n",
    "    \"\"\"\n",
    "    \n",
    "    # This will overwrite the heatmap; useful for changing appearance\n",
    "    # of the plot without recomputing everything\n",
    "    label_keys = []\n",
    "    for label in labels:\n",
    "        if label != -1 and label not in label_keys:\n",
    "            label_keys.append(label)\n",
    "    label_keys.sort()\n",
    "\n",
    "    if closest_matches is not None:\n",
    "\n",
    "        if xlim is not None:\n",
    "            map_end = xlim\n",
    "        else:\n",
    "            map_end = len(poses_series)\n",
    "        heatmap = np.zeros((len(label_keys)*cell_height, map_end), dtype=int)\n",
    "\n",
    "        for l, label in enumerate(labels):\n",
    "            f,p = descriptors[l]\n",
    "            if f >= map_end:\n",
    "                continue\n",
    "            if label >= 0:\n",
    "                label_index = label_keys.index(label)\n",
    "                for r in range(cell_height):\n",
    "                    heatmap[(label_index*cell_height)+r,f] += 2\n",
    "            else:\n",
    "                if (f,p) in closest_matches:\n",
    "                    closest_match = closest_matches[(f,p)]\n",
    "                    for r in range(cell_height):\n",
    "                        heatmap[(closest_match*cell_height)+r,f] += 1 \n",
    "    \n",
    "    fig = plt.figure(figsize=(12,6), constrained_layout=True)\n",
    "    fig.dpi=100\n",
    "    ax = plt.gca()\n",
    "    im = ax.imshow(heatmap, cmap='viridis_r')\n",
    "    print(cell_height/2, ((len(label_keys)*cell_height)+cell_height/2, cell_height))\n",
    "    ax.set_yticks(np.arange(cell_height/2, (len(label_keys)*cell_height)+cell_height/2, cell_height))\n",
    "    ax.set_yticklabels(np.arange(len(label_keys)))\n",
    "    \n",
    "    if video_file is not None:\n",
    "        fps, total_frames = get_video_stats(video_file)\n",
    "        if fps > 0:\n",
    "            ax.set_xticks(np.arange(0,len(poses_series[:map_end]), fps*15))\n",
    "            ax.set_xticklabels([int(poses_series[k]['time']) for k in range(0,map_end,int(round(fps*15)))])\n",
    "            if time_index is not None:\n",
    "                ax.axvline(x=time_index*fps, ymin=0, ymax=len(label_keys)*cell_height,color='r')\n",
    "    fig.tight_layout()\n",
    "    if show:\n",
    "        plt.show()\n",
    "        \n",
    "    return heatmap\n",
    "\n",
    "def compute_pose_distribution(poses_series, labels, descriptors, figure_type='zeroified_figures', cluster_averages=None):\n",
    "    \"\"\" Assign non-clustered poses to clusters (can take a long time) and generate\n",
    "        a clustering timeline heatmap of the pose occurrences.\n",
    "    \"\"\"\n",
    "    \n",
    "    label_keys = []\n",
    "    for label in labels:\n",
    "        if label != -1 and label not in label_keys:\n",
    "            label_keys.append(label)\n",
    "    label_keys.sort()\n",
    "    \n",
    "    heatmap = np.zeros((len(label_keys)*CELL_HEIGHT, len(poses_series)), dtype=int)\n",
    "\n",
    "    closest_matches = {}\n",
    "\n",
    "    for l, label in enumerate(labels):\n",
    "        f,p = descriptors[l]\n",
    "        if label >= 0:\n",
    "            label_index = label_keys.index(label)\n",
    "            for r in range(CELL_HEIGHT):\n",
    "                heatmap[(label_index*CELL_HEIGHT)+r,f] += 2\n",
    "                closest_matches[(f,p)] = label_index\n",
    "        else:\n",
    "            if cluster_averages is None:\n",
    "                continue\n",
    "            print(\"Assigning item\",l,\"of\",len(labels),\"frame\",f,\"pose\",p)\n",
    "            pose_matrix = matrixify_pose(poses_series[f][figure_type][p].data)\n",
    "            if pose_matrix is not None:\n",
    "                match_index = find_nearest_pose(pose_matrix,cluster_averages)\n",
    "                if match_index in label_keys:\n",
    "                    closest_match = label_keys.index(match_index)\n",
    "                    for r in range(CELL_HEIGHT):\n",
    "                        heatmap[(closest_match*CELL_HEIGHT)+r,f] += 1\n",
    "                        closest_matches[(f,p)] = closest_match   \n",
    "    \n",
    "    return heatmap, closest_matches\n",
    "\n",
    "\n",
    "def condense_labels(labels, cluster_map):\n",
    "    \"\"\" Can be used to \"collapse\" clusters of similar poses into meta-clusters\n",
    "    \"\"\"\n",
    "    new_labels = labels\n",
    "    \n",
    "    for l, label in enumerate(labels):\n",
    "        if label != -1 and label in cluster_map:\n",
    "            new_labels[l] = cluster_map[label]\n",
    "\n",
    "    return new_labels\n",
    "\n",
    "\n",
    "def compare_multiple(pose_data, method='distance', figure_type='aligned_figures'):\n",
    "    \"\"\" For multi-dancer videos: Get the mean and standard deviation of inter-pose\n",
    "        similarities for each frame\n",
    "    \"\"\"\n",
    "    frame_means = []\n",
    "    frame_stdevs = []\n",
    "    for f, frame in enumerate(pose_data):\n",
    "        print(\"Processing frame\",f,\"of\",len(pose_data))\n",
    "        frame_similarities = []\n",
    "        for i, figure_i in enumerate(frame[figure_type]):\n",
    "            for j, figure_j in enumerate(frame[figure_type]):\n",
    "                if i < j:\n",
    "                    if method == 'distance':\n",
    "                        mi = get_pose_matrix(frame, i)\n",
    "                        mj = get_pose_matrix(frame, j)\n",
    "                        if mi is None or mj is None:\n",
    "                            similarity = np.nan\n",
    "                        else:\n",
    "                            similarity = mantel(mi, mj)[0]\n",
    "                    else: # method == 'laplacian'\n",
    "                        mi = get_laplacian_matrix(frame, i)\n",
    "                        mj = get_laplacian_matrix(frame, j)\n",
    "                        if mi is None or mj is None:\n",
    "                            similarity = np.nan\n",
    "                        else:\n",
    "                            similarity = 1 - abs(np.subtract(mi.todense(), mj.todense()).sum())\n",
    "                    frame_similarities.append(similarity)\n",
    "    \n",
    "        frame_means.append(np.nanmean(frame_similarities))\n",
    "        frame_stdevs.append(np.nanstd(frame_similarities))\n",
    "\n",
    "    return [frame_means, frame_stdevs]\n",
    "\n",
    "\n",
    "def plot_interpose_similarity(pose_series, frame_means, frame_stdevs, video_file, show=False, min_clip=.2):\n",
    "    \"\"\" For multi-pose videos \"\"\"\n",
    "    fps, total_frames = get_video_stats(video_file)\n",
    "    print(fps,total_frames)\n",
    "    window_length = max(5, int(round(fps/2.0)))\n",
    "\n",
    "    timecodes = []\n",
    "    std_uppers = []\n",
    "    std_lowers = []\n",
    "    total_frames = min(len(pose_series), len(frame_means))\n",
    "    for i, frame in enumerate(pose_series[:total_frames]):\n",
    "        timecodes.append(max(min_clip, frame['time']))\n",
    "        std_uppers.append(max(min_clip, min(1,frame_means[i] + frame_stdevs[i])))\n",
    "        std_lowers.append(max(min_clip, max(0,frame_means[i] - frame_stdevs[i])))\n",
    "\n",
    "    smoothed_means = smooth_series(np.array(frame_means[:total_frames]),window_length)\n",
    "    smoothed_uppers = smooth_series(np.array(std_uppers), window_length)\n",
    "    smoothed_lowers = smooth_series(np.array(std_lowers), window_length)\n",
    "\n",
    "    means_mean = np.nanmean(np.array(frame_means))\n",
    "    smoothed_means_mean = np.nanmean(np.asarray(smoothed_means))\n",
    "    means_stdv = np.nanstd(np.array(frame_means))\n",
    "    smoothed_means_stdv = np.nanstd(np.asarray(smoothed_means))\n",
    "\n",
    "    THRESH = .9\n",
    "\n",
    "    frame_means_array = np.asarray(frame_means)\n",
    "\n",
    "    sim_above_thresh = (frame_means_array > THRESH).sum()\n",
    "\n",
    "    pct_over_thresh = float((np.asarray(frame_means) > THRESH).sum()) / float(len(frame_means))\n",
    "    smoothed_pct_over_thresh = float((np.asarray(smoothed_means) > THRESH).sum()) / float(len(smoothed_means))\n",
    "\n",
    "    if show:\n",
    "        fig = plt.figure(figsize=(12,6), constrained_layout=True)\n",
    "        fig.dpi=100\n",
    "        plt.plot(timecodes, smoothed_means, label=\"mean similarity\")\n",
    "        plt.plot(timecodes, smoothed_uppers, label=\"upper stdev\", linestyle=':')\n",
    "        plt.plot(timecodes, smoothed_lowers, label=\"lower stdev\", linestyle=':')\n",
    "        plt.show()\n",
    "    \n",
    "    return [smoothed_means, smoothed_uppers, smoothed_lowers, timecodes]\n",
    "\n",
    "\n",
    "def average_frame_movements(movement_series, poses_series, show=False, max_clip=3, video_file=None):\n",
    "    \"\"\" Compute the average and stdv of the inter-frame movement for each frame of a\n",
    "        pose sequence with >= 1 dancers.\n",
    "        Assumes movement_series is an array of inter-frame movement values, one for each\n",
    "        detected pose, with missing poses identified via np.nan, as generated from\n",
    "        process_movement_series.\n",
    "    \"\"\"\n",
    "    if len(movement_series) == 0:\n",
    "        print(\"ERROR: empty movement series\")\n",
    "    # Each row should have the same length, so use the first one\n",
    "    total_frames = min(len(movement_series[0]), len(poses_series))\n",
    "    total_poses = len(movement_series)\n",
    "    frame_means = []\n",
    "    upper_stdvs = []\n",
    "    lower_stdvs = []\n",
    "    timecodes = []\n",
    "    for frame in poses_series[:total_frames]:\n",
    "        timecodes.append(frame['time'])\n",
    "    for f in range(total_frames):\n",
    "        frame = []\n",
    "        for p in range(total_poses):\n",
    "            if f < len(movement_series[p]):\n",
    "                frame.append(movement_series[p][f])\n",
    "        if len(frame) == 0 or np.isnan(np.nanmean(frame)):\n",
    "            frame_stdv = 0\n",
    "            upper_stdv = 0\n",
    "            lower_stdv = 0\n",
    "            frame_mean = 0\n",
    "        else:\n",
    "            frame_mean = min(max_clip, np.nanmean(frame))\n",
    "            frame_stdv = np.nanstd(frame)\n",
    "            upper_stdv = min(frame_mean + frame_stdv, max_clip)\n",
    "            lower_stdv = min(frame_mean - frame_stdv, max_clip)\n",
    "        frame_means.append(frame_mean)\n",
    "        upper_stdvs.append(upper_stdv)\n",
    "        lower_stdvs.append(lower_stdv)\n",
    "\n",
    "    if video_file is not None:\n",
    "    \n",
    "        fps, total_frames = get_video_stats(video_file)\n",
    "\n",
    "        mean_mvt = np.nanmean(np.array(frame_means))\n",
    "        mean_stdv = np.nanstd(np.array(frame_means))\n",
    "\n",
    "        mean_mvt_ps = mean_mvt * fps\n",
    "\n",
    "    if show:\n",
    "        fig = plt.figure(figsize=(12,6), constrained_layout=True)\n",
    "        fig.dpi=100\n",
    "        plt.plot(timecodes, frame_means, label=\"mean movement\")\n",
    "        plt.plot(timecodes, upper_stdvs, label=\"upper stdev\", linestyle=':')\n",
    "        plt.plot(timecodes, lower_stdvs, label=\"lower stdev\", linestyle=':')\n",
    "        plt.show()\n",
    "    \n",
    "    return [frame_means, upper_stdvs, lower_stdvs, timecodes]\n",
    "\n",
    "\n",
    "def member_frame_movements(movement_series, poses_series, max_clip=3, show=False, condense=True):\n",
    "    \"\"\" For a multi-dancer video, generate individual inter-frame movement series\n",
    "        for all of the dancers and plot them.\n",
    "    \"\"\"\n",
    "    \n",
    "    print(len(movement_series),\"DANCERS TO CHECK\")\n",
    "    total_frames = min(len(movement_series[0]), len(poses_series))\n",
    "    # Remove series for dancers who never move (due to clipping of sequence)\n",
    "    valid_series = []\n",
    "    for d, dancer in enumerate(movement_series):\n",
    "        if (len(dancer[:total_frames]) == 0) or np.isnan(np.nanmax(dancer[:total_frames])):\n",
    "          print(\"DANCER NEVER MOVES, SKIPPING\")\n",
    "          continue\n",
    "        else:\n",
    "            for v, val in enumerate(dancer[:total_frames]):\n",
    "                if val > max_clip:\n",
    "                    dancer[v] = max_clip\n",
    "            valid_series.append(dancer[:total_frames])\n",
    "    timecodes = []\n",
    "    for frame in poses_series[:total_frames]:\n",
    "        timecodes.append(frame['time'])\n",
    "    valid_array = np.transpose(np.array(valid_series))\n",
    "    \n",
    "    max_frame_figures = 0\n",
    "    for frame in valid_array:\n",
    "        max_frame_figures = max(max_frame_figures,np.count_nonzero(~np.isnan(frame)))\n",
    "\n",
    "    condensed_array = np.zeros((valid_array.shape[0],max_frame_figures),dtype=float)\n",
    "\n",
    "    figures_per_frame = []\n",
    "    for f, frame in enumerate(valid_array):\n",
    "        figures_this_frame = np.count_nonzero(~np.isnan(frame))\n",
    "        max_frame_figures = max(max_frame_figures,figures_this_frame)\n",
    "        figures_per_frame.append(figures_this_frame)\n",
    "\n",
    "        non_nans = np.argwhere(~np.isnan(frame))\n",
    "        for m in range(0,max_frame_figures):\n",
    "            if m < len(non_nans):\n",
    "                condensed_array[f,m] = frame[non_nans[m]]\n",
    "            else:\n",
    "                condensed_array[f,m] = np.nan\n",
    "    \n",
    "    if show:\n",
    "    \n",
    "        fig = plt.figure(figsize=(12,6), constrained_layout=True)\n",
    "        fig.dpi=100\n",
    "        if not condense:\n",
    "            plt.plot(timecodes, valid_array)\n",
    "        else:\n",
    "            plt.plot(timecodes, condensed_array)\n",
    "        plt.show()\n",
    "\n",
    "    if not condense:\n",
    "        return valid_array\n",
    "    else:\n",
    "        return condensed_array\n",
    "\n",
    "def compare_sequences_pairwise(seq1, seq2, figure_type='figures'):\n",
    "    def pose_match_score(matrix1, matrix2):\n",
    "        return mantel(matrix1, matrix2)[0]\n",
    "\n",
    "    matrix_seq1 = [get_pose_matrix(frame, 0, figure_type) for frame in seq1]\n",
    "    matrix_seq2 = [get_pose_matrix(frame, 0, figure_type) for frame in seq2]\n",
    "\n",
    "    pairwise2.align.localcx(matrix_seq1, matrix_seq2, pose_match_score, gap_char=[])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.10 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "c6e4e9f98eb68ad3b7c296f83d20e6de614cb42e90992a65aa266555a3137d0d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}