# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.8.1 64-bit (''env'': venv)'
#     language: python
#     name: python38164bitenvvenv181da0e8f8774db09a4bf241dbbabaf9
# ---

# # Results for 'Evaluating the LENA System for Korean'

# + tags=[]
import os
import sys
import json

from evaluation import *

import pyannote.core.notebook
from pyannote.metrics.diarization import DiarizationErrorRate, IdentificationErrorRate
from pyannote.core import Segment, Annotation

import textgrid
import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.metrics import precision_recall_fscore_support, confusion_matrix, classification_report

from collections import Counter

# %matplotlib inline
sns.set()
# -

# ## Inter-rater reliability

# +
clips = (1, 2, 3, 4, 5, 27, 53, 54)
irr_results = pd.DataFrame(columns=['clip', 'identification error rate', 'missed detection', 'false alarm', 'confusion', 'correct', 'total', 'awc', 'cvc', 'ctc'])

# The same mappings are applied as in the main IER analysis
with open('mappings/ours_ier.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

for clip in clips:
    filename_1 = 'data/textgrids/Clip{}.TextGrid'.format(clip) # original
    filename_2 = 'data/textgrids_irr/Clips{}_suhan_CT.TextGrid'.format(clip)

    transcript_dict_1 = textgrid_to_dict(filename_1)
    # transcript_annot_1 = dict_to_annotation(transcript_dict_1, silence_class='Silence')
    transcript_annot_1 = dict_to_annotation(
            remap(transcript_dict_1, human_to_ground_mappings), 
            silence_class='Silence')

    transcript_dict_2 = textgrid_to_dict(filename_2)
    # transcript_annot_2 = dict_to_annotation(transcript_dict_2, silence_class='Silence')
    transcript_annot_2 = dict_to_annotation(
            remap(transcript_dict_2, human_to_ground_mappings), 
            silence_class='Silence')

    included_in_awc = ['Female', 'Male', 'Female2', 'Male2']
    exclude_from_vocalizations = ('F', 'V')
    include_as_turn = ('CA', 'AC', 'CO', 'OC')

    awc = count_word_whitespace(None, path=filename_2, included_tiers=included_in_awc)
    cvc = count_child_vocalizations(filename_2, excluded=exclude_from_vocalizations)
    ctc = count_conversational_turns(filename_2, included=include_as_turn)
    ier = IdentificationErrorRate(collar=0.0, skip_overlap=True)

    results = ier(transcript_annot_1, transcript_annot_2, detailed=True)
    results['clip'] = clip
    results['awc'] = awc
    results['cvc'] = cvc
    results['ctc'] = ctc
    
    irr_results = irr_results.append(results, ignore_index=True)
# -

irr_results


def plot_confusion(confusion, norm=None, save=False):
    """ Plots a confusion matrix displaying one of precision, recall, and raw counts.
    """
    if norm == 'recall':
        normed_confusion = confusion.astype('float') / confusion.sum(axis=1)[:, np.newaxis]
        fmt = '.0%'
    elif norm == 'precision':
        normed_confusion = confusion.astype('float') / confusion.sum(axis=0)[np.newaxis, :]
        fmt = '.0%'
    else:
        normed_confusion = confusion
        fmt = 'd'
        
    normed_confusion = normed_confusion.fillna(0)

    fig = plt.figure(figsize=(8, 6))

    if norm == 'recall':
        plt.title('Recall', fontweight='bold', fontsize=15, pad=12)
        cmap = 'Purples'
    elif norm == 'precision':
        plt.title('Precision', fontweight='bold', fontsize=15, pad=12)
        cmap = 'Greens'
    else:
        plt.title('Confusion matrix')
        cmap = 'Reds'

    ax = sns.heatmap(normed_confusion, annot=True, fmt=fmt,  cbar=False, cmap=cmap)
    # ax = sns.heatmap(confusion, annot=True, xticklabels=labels, yticklabels=labels, cbar=False, cmap="YlGnBu")
    # ax.set(xlabel='LENA', ylabel='Human')
    ax.set_ylabel('Human', labelpad=12, fontweight='bold')
    plt.yticks(np.arange(len(normed_confusion.index))+0.5, rotation=0, fontsize="10", va="center")
    ax.set_xlabel('LENA', labelpad=12, fontweight='bold')

    # ax.set_ylim([0,2])

    plt.show()

    if save:
        fig.savefig('results/{}.pdf'.format(norm), dpi=600, bbox_inches='tight')

    return normed_confusion



# ## Confusion matrices for main analysis

# + tags=[]
with open('mappings/ours_confusion.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

y_lena, y_human = clip_to_frames_all(lena_to_ground_mappings, human_to_ground_mappings, overlap_classes)

categorical_human = pd.Categorical(y_human, categories=human_order)
categorical_lena = pd.Categorical(y_lena, categories=lena_order)

cross_tabulation = pd.crosstab(categorical_human, categorical_lena)

# cross_tabulation = cross_tabulation.rename(index={'UnspecifiedOverlap': 'Other'})
# -

# Computed results
cross_tabulation = pd.read_csv('results/confusion_main.csv', index_col='row_0')
cross_tabulation

plot_confusion(cross_tabulation, 'precision', True)

plot_confusion(cross_tabulation, 'recall', True)

# + tags=[]
print(classification_report(y_human, y_lena, labels=human_order))
# -

# ## Supplementary confusion

# + tags=[]
with open('mappings/cristia2020.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

y_lena, y_human = clip_to_frames_all(lena_to_ground_mappings, human_to_ground_mappings, overlap_classes)

categorical_human = pd.Categorical(y_human, categories=human_order)
categorical_lena = pd.Categorical(y_lena, categories=lena_order)

cross_tabulation = pd.crosstab(categorical_human, categorical_lena)

# cross_tabulation = cross_tabulation.rename(index={'UnspecifiedOverlap': 'Other'})
# -

cross_tabulation.to_csv('results/confusion_supplementary.csv', index=True)

plot_confusion(cross_tabulation, 'recall', True)

plot_confusion(cross_tabulation, 'precision', True)

# ## Identification Error Rate--Main

with open('mappings/ours_ier.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

# + tags=[]
ier = IdentificationErrorRate(collar=0.0, skip_overlap=True)

df = pd.DataFrame(columns=['clip', 'identification error rate', 'missed detection', 'false alarm', 'confusion', 'correct', 'total'])

for i in range(1, 61):
    print('Calculating IER for clip', i)
    human_annot, lena_annot = clip_to_annotations(i, lena_to_ground_mappings, human_to_ground_mappings)
    results = ier(human_annot, lena_annot, detailed=True)
    results['clip'] = i
    df = df.append(results, ignore_index=True)
# -

df

# ### IER calculated manually using sampled frames
# (This should be nearly identical to the above)

# + tags=[]
frame_ier = []
speech_tiers = ['Female', 'Male', 'Child', 'OtherChild']
for i in range(1, 61):
    f_l, f_h = clip_to_frames_single(i, lena_to_ground_mappings, human_to_ground_mappings, overlap_classes)

    false_alarms, misses, confusions, total = calculate_frame_ier(f_h, f_l, speech_tiers, True)
    d = (false_alarms + misses + confusions) / total 
    frame_ier.append(d)

for error in frame_ier:
    print(error)
# -

# # IER--Supplementary

# ## Including TV

# + tags=[]
with open('mappings/ier_include_tv.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

ier = IdentificationErrorRate(collar=0.0, skip_overlap=True)

df = pd.DataFrame(columns=['clip', 'identification error rate', 'missed detection', 'false alarm', 'confusion', 'correct', 'total'])

for i in range(1, 61):
    print('Calculating IER for clip', i)
    human_annot, lena_annot = clip_to_annotations(i, lena_to_ground_mappings, human_to_ground_mappings)
    results = ier(human_annot, lena_annot, detailed=True)
    results['clip'] = i
    df = df.append(results, ignore_index=True)

df

# + tags=[]
# Verify with frame-level IER
frame_ier = []
speech_tiers = ['Female', 'Male', 'Child', 'OtherChild', 'TV']
for i in range(1, 61):
    f_l, f_h = clip_to_frames_single(i, lena_to_ground_mappings, human_to_ground_mappings, overlap_classes)

    false_alarms, misses, confusions, total = calculate_frame_ier(f_h, f_l, speech_tiers, True)
    d = (false_alarms + misses + confusions) / total 
    frame_ier.append(d)

for error in frame_ier:
    print(error)
# -

# ## Including TV and Overlap

# + tags=[]
with open('mappings/ier_include_tv_overlap.json') as f:
    mappings = json.load(f)
    overlap_classes = mappings['overlap']
    lena_to_ground_mappings = mappings['lena_to_ground_mappings']
    human_to_ground_mappings = mappings['human_to_ground_mappings']
    human_order = mappings['human_order']
    lena_order = mappings['lena_order']

frame_ier = []
speech_tiers = ['Female', 'Male', 'Child', 'OtherChild', 'TV', 'Overlap']

frame_confusions = 0 # need totals
frame_misses = 0
frame_false_alarms = 0
frame_totals = 0

for i in range(1, 61):
    f_l, f_h = clip_to_frames_single(i, lena_to_ground_mappings, human_to_ground_mappings, overlap_classes)

    false_alarms, misses, confusions, total = calculate_frame_ier(f_h, f_l, speech_tiers, skip_overlap=False)
    frame_false_alarms += false_alarms
    frame_misses += misses
    frame_confusions += confusions
    frame_totals += total

print('false alarm:', frame_false_alarms / frame_totals)
print('confusion', frame_confusions / frame_totals)
print('miss', frame_misses / frame_totals)
print('ier', (frame_confusions + frame_false_alarms + frame_misses) / frame_totals)

# -

# ## Conversational Turn Count, Child Vocalization Count

# +
df = pd.read_csv(METADATA_PATH, index_col='ClipNumber')

vocalization_lena = df['CVC'].values
turn_lena = df['CTC'].values

exclude_from_vocalizations = ('F', 'V')
include_as_turn = ('CA', 'AC', 'CO', 'OC')

turn_human = []
vocalization_human = []

for i in range(1, 61):
    filename = 'Clip{}.TextGrid'.format(i)
    filepath = os.path.join(TEXTGRID_PATH, filename)
    turn = count_conversational_turns(filepath, included=include_as_turn)
    vocalization = count_child_vocalizations(filepath, excluded=exclude_from_vocalizations)
    turn_human.append(turn)
    vocalization_human.append(vocalization)
    
turn_human = np.array(turn_human)
vocalization_human = np.array(vocalization_human)
        

# -

turn_human

vocalization_human

# ## Adult Word Count

word_counts = []
included = ['Female', 'Male', 'Female2', 'Male2']
for i in range(1, 61):
    count = count_word_whitespace(i, included)
    word_counts.append(count)

word_counts

# ## Morpheme Count

morph_counts = []
included = ['Female', 'Male', 'Female2', 'Male2']
for i in range(1, 61):
    count = count_morphemes_mecab(i, included)
    morph_counts.append(count)

morph_counts


# +
def plot_all_clip_comparisons():
    for i in range(1, 61):
        human_annot, lena_annot = clip_to_annotations(i, raw_lena_mappings, raw_textgrid_mappings)
        # Plot ground truth and lena side by side
        # create a figure with 6 rows with matplotlib
        nrows = 2
        fig, ax = plt.subplots(nrows=nrows, ncols=1)
        fig.set_figwidth(20)
        fig.set_figheight(nrows * 2)
        # 1st row: reference annotation
        pyannote.core.notebook.plot_annotation(human_annot, ax=ax[1], time=True, legend=True)

        # 2nd row: SAD raw scores
        pyannote.core.notebook.plot_annotation(lena_annot, ax=ax[0], time=False, legend=True)

        ax[0].set_title('Clip {}'.format(i), loc='right')

        fig.savefig('/data/speech/ko_results/raw_vis_{}'.format(i), transparent=False, bbox_inches='tight')
#         plt.show()
        fig.clear()
                   
plot_all_clip_comparisons()

StarcoderdataJupyterScriptsDedupFiltered