SPAM annotation correction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from joblib import Parallel, delayed\n",
    "import pandas as pd\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import librosa\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.2.1\n"
     ]
    }
   ],
   "source": [
    "import jams\n",
    "print(jams.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Set a padding tolerance of 100ms\n",
    "TOLERANCE = 0.1\n",
    "\n",
    "# Set an alignment tolerance of 3s\n",
    "ALIGN_TOLERANCE = 3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def merge_annotations(upper=None, lower=None):\n",
    "    \n",
    "    ann = jams.Annotation(namespace='multi_segment',\n",
    "                          annotation_metadata=upper.annotation_metadata)\n",
    "    \n",
    "    try:\n",
    "        ann.sandbox.update(**upper.sandbox)\n",
    "    except AttributeError:\n",
    "        pass\n",
    "    \n",
    "    try:\n",
    "        ann.sandbox.update(**lower.sandbox)\n",
    "    except AttributeError:\n",
    "        pass\n",
    "    \n",
    "    # Add the uppers\n",
    "    d_upper = upper.data.copy()\n",
    "    d_upper['value'] = [{'label': _, 'level': 0} for _ in d_upper['value']]\n",
    "    \n",
    "    # Add the lowers\n",
    "    d_lower = lower.data.copy()\n",
    "    d_lower['value'] = [{'label': _, 'level': 1} for _ in d_lower['value']]\n",
    "    \n",
    "    ann.data = jams.JamsFrame.from_dataframe(pd.concat([d_upper, d_lower], ignore_index=True))\n",
    "    \n",
    "    return ann\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def span_segment(ann=None, duration=None):\n",
    "    \n",
    "    if ann.namespace == 'segment_salami_upper':\n",
    "        pre, post = 'YYYYY', 'ZZZZZ'\n",
    "    elif ann.namespace == 'segment_salami_lower':\n",
    "        pre, post = 'yyyyy', 'zzzzz'\n",
    "    else:\n",
    "        return\n",
    "    \n",
    "    frame = ann.data\n",
    "    \n",
    "    # Give a one-second buffer here\n",
    "    if (frame['time'] > pd.Timedelta(duration + 1.0, unit='s')).any():\n",
    "        raise RuntimeError('track length exceeded in observation: {0:3g}/{1:3g}'.format(\n",
    "                frame['time'].max().total_seconds(),\n",
    "                duration))\n",
    "        \n",
    "    # Drop any annotations where time > duration\n",
    "    frame = jams.JamsFrame(frame[frame['time'] <= pd.Timedelta(duration, unit='s')])\n",
    "    \n",
    "    idx = frame['time'].argmin()\n",
    "    min_time = frame['time'][idx]\n",
    "    \n",
    "    # If the minimum time is close enough to zero, just clamp it\n",
    "    if pd.Timedelta(0) < min_time:\n",
    "        if min_time <= pd.Timedelta(TOLERANCE, unit='s') or frame['value'][idx].lower() == 'silence':\n",
    "            frame.loc[idx, 'time'] = pd.Timedelta(0)\n",
    "            \n",
    "        else:\n",
    "            # We need to pad a new segment here\n",
    "            \n",
    "            frame.add_observation(time=0,\n",
    "                                  duration=min_time.total_seconds(),\n",
    "                                  value=pre, confidence=0)\n",
    "            \n",
    "        \n",
    "    idx = (frame['time'] + frame['duration']).argmax()\n",
    "    max_time = frame['time'][idx] + frame['duration'][idx]\n",
    "    \n",
    "    if max_time != pd.Timedelta(duration, unit='s'):\n",
    "        if pd.Timedelta(duration - TOLERANCE, unit='s') <= max_time or frame['value'][idx].lower() == 'silence':\n",
    "            frame.loc[idx, 'duration'] = pd.Timedelta(duration - frame['time'][idx].total_seconds(), unit='s')\n",
    "            \n",
    "        else:\n",
    "            # Pad out a new segment here\n",
    "            frame.add_observation(time=max_time.total_seconds(),\n",
    "                                  duration=duration - max_time.total_seconds(),\n",
    "                                  value=post, confidence=0)\n",
    "    \n",
    "    # Sort the rows\n",
    "    frame.sort_values('time', inplace=True)\n",
    "    \n",
    "    ann.data = frame\n",
    "    \n",
    "def align_annotations(upper=None, lower=None):\n",
    "    \n",
    "    # Every upper-level segment should have a corresponding lower-level segment within TOLERANCE\n",
    "    \n",
    "    i_upper, l_upper = upper.data.to_interval_values()\n",
    "    i_lower, l_lower = lower.data.to_interval_values()\n",
    "    \n",
    "    starts, ends = i_upper[:, 0], i_upper[:, 1]\n",
    "    t_times = np.unique(np.ravel(i_lower))\n",
    "    start_match = librosa.util.match_events(starts, t_times)\n",
    "    end_match = librosa.util.match_events(ends, t_times)\n",
    "    \n",
    "    starts_adj = t_times[start_match]\n",
    "    ends_adj = t_times[end_match]\n",
    "    \n",
    "    assert np.all(np.abs(starts_adj - starts) <= ALIGN_TOLERANCE), np.max(np.abs(starts - starts_adj))\n",
    "    \n",
    "    new_df = jams.JamsFrame()\n",
    "        \n",
    "    # Only add intervals with positive duration\n",
    "    for s, t, l, c in zip(starts_adj, ends_adj, upper.data.value, upper.data.confidence):\n",
    "        if t > s:\n",
    "            new_df.add_observation(time=s, duration=t-s, value=l, confidence=c)\n",
    "        \n",
    "    upper.data = new_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def fix_jams(jamsfile):\n",
    "    \n",
    "    J = jams.load(jamsfile)\n",
    "    \n",
    "    ann_dict = defaultdict(dict)\n",
    "    \n",
    "    for upper in J.search(namespace='segment_salami_upper'):\n",
    "        name = upper.annotation_metadata.annotator.name\n",
    "        try:\n",
    "            \n",
    "            span_segment(upper, duration=J.file_metadata.duration)\n",
    "            ann_dict[name]['upper'] = upper\n",
    "        except RuntimeError as exc:\n",
    "            print(exc, jamsfile, name, 'upper\\n')\n",
    "\n",
    "    for lower in J.search(namespace='segment_salami_lower'):\n",
    "        name = lower.annotation_metadata.annotator.name\n",
    "        try:\n",
    "            \n",
    "            span_segment(lower, duration=J.file_metadata.duration)\n",
    "            ann_dict[name]['lower'] = lower\n",
    "    \n",
    "        except RuntimeError as exc:\n",
    "            print(exc, jamsfile, name, 'lower\\n')\n",
    "\n",
    "    for key in ann_dict:\n",
    "        if len(ann_dict[key]) != 2:\n",
    "            print('Align and merge failed: {}/{}'.format(jamsfile, key))\n",
    "            continue\n",
    "        try:\n",
    "            align_annotations(**ann_dict[key])\n",
    "            J.annotations.append(merge_annotations(**ann_dict[key]))\n",
    "            pass\n",
    "        except AssertionError as exc:\n",
    "            print(exc, jamsfile, key)\n",
    "\n",
    "    return J"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "files = jams.util.find_with_extension('/home/bmcfee/git/msaf-data/SPAM/references/', 'jams')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27.028027 /home/bmcfee/git/msaf-data/SPAM/references/SALAMI_108.jams Shuli Tang\n",
      "track length exceeded in observation: 401.914/297.874 /home/bmcfee/git/msaf-data/SPAM/references/SALAMI_114.jams Shuli Tang upper\n",
      "\n",
      "Align and merge failed: /home/bmcfee/git/msaf-data/SPAM/references/SALAMI_114.jams/Shuli Tang\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for jfn in tqdm(files):\n",
    "    J = fix_jams(jfn)\n",
    "    J.save(jfn)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.5",
   "language": "python",
   "name": "python3.5"
  },
  "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.5.0"
  },
  "widgets": {
   "state": {
    "04a52b1a10d04e199eac3b3687b6ee38": {
     "views": []
    },
    "0ce201e2577c4d6a9302c8103baa960e": {
     "views": []
    },
    "1022201d3ce84cdea296a50ed5050818": {
     "views": []
    },
    "14e4a650dfad46bfb4a8f3534a03f6f0": {
     "views": []
    },
    "157f106726d9425f895b585abe584f2f": {
     "views": []
    },
    "1771d70dab4247dea81f15eb08d2d81c": {
     "views": []
    },
    "187c0c4890a04ead997a1fedf502f3d4": {
     "views": []
    },
    "1b265a2610214deba55ad48f77377a36": {
     "views": []
    },
    "1fd6664860fc4ac986e649e432c37d25": {
     "views": []
    },
    "1fe38e6ab9b9477ca1d5e2bfb4b77213": {
     "views": []
    },
    "216c33cfd35a4828aa0ab81accdda254": {
     "views": []
    },
    "27d0d09a0b854d0883cd8def3d72d712": {
     "views": []
    },
    "2d48fbd453984519b9b13518c25952cf": {
     "views": []
    },
    "465a25f4db9a4bfaa984c9bcfa43e3af": {
     "views": []
    },
    "5822cb6580074690a558009ac82c11a6": {
     "views": []
    },
    "5841535b125c4407b433506a02c37a79": {
     "views": []
    },
    "613348f310bf43748bcaa5302df805bc": {
     "views": []
    },
    "6999ab34750a4f77bdb8c428c54964f0": {
     "views": []
    },
    "73f217d2cd6a4deb973c0a53cc42d5da": {
     "views": []
    },
    "777dc4b51a1547c2b5d5806cef8a5838": {
     "views": []
    },
    "78f6080393ed49429322c85481c1e6d2": {
     "views": []
    },
    "7f1b9bb3a4094fc496726fe175a09724": {
     "views": []
    },
    "8212de3ef8754b72b5630dc4486cab16": {
     "views": []
    },
    "882ce93002b045918bbac41e2fb82f22": {
     "views": []
    },
    "9340168f988c435f949f4b1eb8b16a86": {
     "views": []
    },
    "96e23e80af51429b92dd0149c70eef4d": {
     "views": []
    },
    "ab162c8bc22b49119fb37911f2313c17": {
     "views": []
    },
    "ac371daca40645619afc3f180619d4ea": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "ac8e1a89062a4175a7ca2790b785b43f": {
     "views": []
    },
    "b0bbf56bd33545ad9baf9c910a86ae8b": {
     "views": []
    },
    "b17bdd591b98486aaf6a7e5b0130625b": {
     "views": []
    },
    "b2d911db28ca489398a15df6cb697fe2": {
     "views": []
    },
    "bb0a4fe4036e485082a709df1c739286": {
     "views": []
    },
    "bb18b737e1d44004bf3bd4448f468f00": {
     "views": []
    },
    "c202989aee1d42f6bb15cf9f1cfcec47": {
     "views": []
    },
    "c28f180a8ea744d19cde88e908f117a1": {
     "views": []
    },
    "c49a52f49cfb483ca3a60114c0e2cdd9": {
     "views": []
    },
    "c91430292aa9498da804c639b284035e": {
     "views": []
    },
    "cc0559d60fef4e69a8c5846d9c79de1b": {
     "views": []
    },
    "e4185f32cfea4777b0d9da948c638679": {
     "views": []
    },
    "ef8531b7d52747e795caf703e5cbe80e": {
     "views": []
    },
    "fb7ddf57d53645a887f1bb238ecfe722": {
     "views": []
    }
   },
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}