Concatenate short segments

hirofumi0810 · hirofumi0810 · commit 827c90bff707 · 2021-04-23T00:27:00.000+09:00
diff --git a/egs/iwslt21/asr1/local/merge_short_segments.py b/egs/iwslt21/asr1/local/merge_short_segments.py
@@ -0,0 +1,121 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Kyoto University (Hirofumi Inaguma)
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Merge adjacent utterances."""
+
+
+import argparse
+import codecs
+from collections import deque
+
+parser = argparse.ArgumentParser()
+parser.add_argument("segments", type=str, help="path to segment file")
+
+parser.add_argument("output_segments", type=str, help="path to output segment file")
+parser.add_argument("output_utt2spk", type=str, help="path to output utt2spk file")
+parser.add_argument("output_spk2utt", type=str, help="path to output spk2utt file")
+
+parser.add_argument("--min_interval", type=int, default=200, help="")
+parser.add_argument(
+    "--max_duration", type=int, default=1500, help="maximum duration [frame]"
+)
+parser.add_argument(
+    "--delimiter", type=str, default="_", help="delimiter on utt_id start_time end_time"
+)
+args = parser.parse_args()
+
+
+def merge(segments, segments_dict):
+
+    while True:
+        num_merge = 0
+        new_segments = deque([])
+        utt_id_prev, start_prev, end_prev = segments.popleft()
+        utt_ids_merged = utt_id_prev
+        for utt_ids, start, end in segments:
+            interval = start - end_prev
+            duration = end - start_prev
+            if interval < args.min_interval and duration < args.max_duration:
+                # merge
+                end_prev = end
+                utt_ids_merged.extend(utt_ids)
+                num_merge += 1
+            else:
+                new_segments.append((utt_ids_merged, start_prev, end_prev))
+
+                start_prev = start
+                end_prev = end
+                utt_ids_merged = utt_ids
+
+        # for last segments
+        new_segments.append((utt_ids_merged, start_prev, end))
+        segments = new_segments
+
+        if num_merge == 0:
+            break
+
+    delimiter = args.delimiter
+    for utt_ids, _, _ in segments:
+        spk = utt_ids[0].split(delimiter)[0]
+        s = utt_ids[0].split(delimiter)[1]
+        e = utt_ids[-1].split(delimiter)[2]
+        new_utt_id = "%s" % (spk + delimiter + s + delimiter + e)
+
+        segments_dict[new_utt_id] = (
+            segments_dict[utt_ids[0]][0],
+            segments_dict[utt_ids[-1]][1],
+        )
+
+        if len(utt_ids) > 1:
+            for utt_id in utt_ids:
+                del segments_dict[utt_id]
+
+    return segments_dict
+
+
+def main():
+    segments_dict = {}
+    with codecs.open(args.segments, "r", encoding="utf-8") as f:
+        for line in f:
+            utt_id, spk, start, end = line.strip().split()
+            segments_dict[utt_id] = (start, end)
+
+    segments_spk = deque([])
+    with codecs.open(args.segments, "r", encoding="utf-8") as f:
+        spk_prev = None
+        for line in f:
+            utt_id, spk, start, end = line.strip().split()
+            start = float(start) * 100  # per 10ms
+            end = float(end) * 100  # per 10ms
+            if spk_prev is not None and spk != spk_prev:
+                segments_dict = merge(segments_spk, segments_dict)
+                segments_spk = deque([])  # reset
+            segments_spk.append(([utt_id], start, end))
+            spk_prev = spk
+
+    with codecs.open(args.output_segments, "w", encoding="utf-8") as f:
+        for utt_id, (start, end) in sorted(segments_dict.items(), key=lambda x: x[0]):
+            spk = utt_id.split(args.delimiter)[0]
+            f.write("%s %s %s %s\n" % (utt_id, spk, start, end))
+
+    spk2utt_dict = {}
+    with codecs.open(args.output_utt2spk, "w", encoding="utf-8") as f:
+        for utt_id, ref in sorted(segments_dict.items(), key=lambda x: x[0]):
+            spk = utt_id.split("_")[0]
+            f.write("%s %s\n" % (utt_id, spk))
+
+            if spk not in spk2utt_dict:
+                spk2utt_dict[spk] = [utt_id]
+            else:
+                spk2utt_dict[spk] += [utt_id]
+
+    with codecs.open(args.output_spk2utt, "w", encoding="utf-8") as f:
+        for spk, utt_ids in sorted(spk2utt_dict.items(), key=lambda x: x[0]):
+            f.write("%s %s\n" % (spk, " ".join(utt_ids)))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/egs/iwslt21/st1/run.sh b/egs/iwslt21/st1/run.sh
@@ -50,6 +50,10 @@ tgt_case=tc
 remove_nonverbal=true  # remove non-verbal labels such as "( Applaus )"
 # NOTE: IWSLT community accepts this setting and therefore we use this by default
 
+# iwslt segmentation related
+max_interval=200
+max_duration=1500
+
 # bpemode (unigram or bpe)
 nbpe=16000
 bpemode=bpe
@@ -69,16 +73,17 @@ set -o pipefail
 mustc_dir=../../must_c
 mustc_v2_dir=../../must_c_v2
 stted_dir=../../iwslt18
+
 # test data directory
-iwslt_test_data=/n/rd8/iwslt18
+iwslt_test_data_dir=/n/rd8/iwslt18
 
 train_set=train.de
 train_dev=dev.de
 trans_subset="et_mustc_dev_org.de et_mustc_tst-COMMON.de et_mustc_tst-HE.de"
 trans_set="et_mustc_dev_org.de et_mustc_tst-COMMON.de et_mustc_tst-HE.de \
-           et_mustcv2_dev_org.de et_mustcv2_tst-COMMON.de et_mustcv2_tst-HE.de \
-           et_stted_dev2010.de et_stted_tst2010.de et_stted_tst2013.de et_stted_tst2014.de et_stted_tst2015.de \
-           et_stted_tst2018.de et_stted_tst2019.de"
+           et_mustcv2_dev_org.de et_mustcv2_tst-COMMON.de et_mustcv2_tst-HE.de"
+iwslt_test_set="et_stted_dev2010.de et_stted_tst2010.de et_stted_tst2013.de et_stted_tst2014.de et_stted_tst2015.de \
+                et_stted_tst2018.de et_stted_tst2019.de"
 
 if [ ${stage} -le 0 ] && [ ${stop_stage} -ge 0 ]; then
     ### Task dependent. You have to make data the following preparation part by yourself.
@@ -146,9 +151,9 @@ if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
         steps/make_fbank_pitch.sh --cmd "$train_cmd" --nj 32 --write_utt2num_frames true \
             data/${x} exp/make_fbank/${x} ${fbankdir}
         utils/fix_data_dir.sh --utt_extra_files "text.tc text.lc text.lc.rm" data/${x}
+        rm data/${x}/segments
+        rm data/${x}/wav.scp
     done
-    rm data/*/segments
-    rm data/*/wav.scp
 
     for lang in en de; do
         utils/combine_data.sh --extra_files "text.tc text.lc text.lc.rm" data/train.${lang} data/tr_mustc.${lang} data/tr_mustcv2.${lang} data/tr_stted.${lang}
@@ -184,6 +189,32 @@ if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
         dump.sh --cmd "$train_cmd" --nj 32 --do_delta $do_delta \
             data/${x}/feats.scp data/${train_set}/cmvn.ark exp/dump_feats/trans/${x} ${feat_trans_dir}
     done
+
+    # concatenate short segments
+    for x in ${iwslt_test_set}; do
+        output_dir=${x}_merge${max_interval}_duration${max_duration}
+        rm -rf data/${output_dir}
+        cp -rf data/${x} data/${output_dir}
+        rm data/${output_dir}/utt2num_frames
+
+        local/merge_short_segments.py \
+            data/${x}/segments \
+            data/${output_dir}/segments \
+            data/${output_dir}/utt2spk \
+            data/${output_dir}/spk2utt \
+            --min_interval ${max_interval} \
+            --max_duration ${max_duration} \
+            --delimiter "_" || exit 1;
+
+        # Generate the fbank features; by default 80-dimensional fbanks with pitch on each frame
+        steps/make_fbank_pitch.sh --cmd "$train_cmd" --nj 32 --write_utt2num_frames true \
+            data/${output_dir} exp/make_fbank/${output_dir} ${fbankdir}
+        utils/fix_data_dir.sh --utt_extra_files "text.tc text.lc text.lc.rm" data/${output_dir}
+
+        feat_trans_dir=${dumpdir}/${output_dir}/delta${do_delta}; mkdir -p ${feat_trans_dir}
+        dump.sh --cmd "$train_cmd" --nj 32 --do_delta $do_delta \
+            data/${output_dir}/feats.scp data/${train_set}/cmvn.ark exp/dump_feats/trans/${output_dir} ${feat_trans_dir}
+    done
 fi
 
 dict=data/lang_1spm/${train_set}_${bpemode}${nbpe}_units_${tgt_case}.txt
@@ -212,22 +243,20 @@ if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
     echo "make json files"
     data2json.sh --nj 16 --feat ${feat_tr_dir}/feats.scp --text data/${train_set}/text.${tgt_case} --bpecode ${bpemodel}.model --lang "de" \
         data/${train_set} ${dict} > ${feat_tr_dir}/data_${bpemode}${nbpe}.${src_case}_${tgt_case}.json
-    for x in ${train_dev} ${trans_set}; do
-        feat_trans_dir=${dumpdir}/${x}/delta${do_delta}
+    for x in ${train_dev} ${trans_set} ${iwslt_test_set}; do
         if [[ ${x} = *tst20* ]] || [[ ${x} = *dev20* ]]; then
+            feat_trans_dir=${dumpdir}/${x}_merge${max_interval}_duration${max_duration}/delta${do_delta}
             local/data2json.sh --feat ${feat_trans_dir}/feats.scp --no_text true \
-                data/${x} ${dict} > ${feat_trans_dir}/data_${bpemode}${nbpe}.${src_case}_${tgt_case}.json
+                data/${x}_merge${max_interval}_duration${max_duration} ${dict} > ${feat_trans_dir}/data_${bpemode}${nbpe}.${src_case}_${tgt_case}.json
         else
+            feat_trans_dir=${dumpdir}/${x}/delta${do_delta}
             data2json.sh --feat ${feat_trans_dir}/feats.scp --text data/${x}/text.${tgt_case} --bpecode ${bpemodel}.model --lang "de" \
                 data/${x} ${dict} > ${feat_trans_dir}/data_${bpemode}${nbpe}.${src_case}_${tgt_case}.json
         fi
     done
 
     # update json (add source references)
     for x in ${train_set} ${train_dev} ${trans_set}; do
-        if [[ ${x} = *tst20* ]] || [[ ${x} = *dev20* ]]; then
-            continue
-        fi
         feat_dir=${dumpdir}/${x}/delta${do_delta}
         data_dir=data/$(echo ${x} | cut -f 1 -d ".").en
         update_json.sh --text ${data_dir}/text.${src_case} --bpecode ${bpemodel}.model \
@@ -297,6 +326,9 @@ if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
     pids=() # initialize pids
     for x in ${trans_subset}; do
     (
+        if [[ ${x} = *tst20* ]] || [[ ${x} = *dev20* ]]; then
+            x=${x}_merge${max_interval}_duration${max_duration}
+        fi
         decode_dir=decode_${x}_$(basename ${decode_config%.*})
         feat_trans_dir=${dumpdir}/${x}/delta${do_delta}
 
@@ -317,7 +349,7 @@ if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
             set=$(echo ${x} | cut -f 1 -d "." | cut -f 3 -d "_")
             local/score_bleu_reseg.sh --case ${tgt_case} --bpe ${nbpe} --bpemodel ${bpemodel}.model \
                 --remove_nonverbal ${remove_nonverbal} \
-                ${expdir}/${decode_dir} ${dict} ${iwslt_test_data} ${set}
+                ${expdir}/${decode_dir} ${dict} ${iwslt_test_data_dir} ${set}
         else
             score_bleu.sh --case ${tgt_case} --bpe ${nbpe} --bpemodel ${bpemodel}.model \
                 --remove_nonverbal ${remove_nonverbal} \
