Blame view

egs/wsj/s5/utils/ctm/resolve_ctm_overlaps.py 11.6 KB
8dcb6dfcb   Yannick Estève   first commit
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
  #! /usr/bin/env python
  
  # Copyright 2014  Johns Hopkins University (Authors: Daniel Povey)
  #           2014  Vijayaditya Peddinti
  #           2016  Vimal Manohar
  # Apache 2.0.
  
  """
  Script to combine ctms with overlapping segments.
  The current approach is very simple. It ignores the words,
  which are hypothesized in the half of the overlapped region
  that is closer to the utterance boundary.
  So if there are two segments
  in the region 0s to 30s and 25s to 55s, with overlap of 5s,
  the last 2.5s of the first utterance i.e. from 27.5s to 30s is truncated
  and the first 2.5s of the second utterance i.e. from 25s to 27.s is truncated.
  """
  
  from __future__ import print_function
  from __future__ import division
  import argparse
  import collections
  import logging
  
  from collections import defaultdict
  
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)
  handler = logging.StreamHandler()
  handler.setLevel(logging.INFO)
  formatter = logging.Formatter(
      '%(asctime)s [%(pathname)s:%(lineno)s - '
      '%(funcName)s - %(levelname)s ] %(message)s')
  handler.setFormatter(formatter)
  logger.addHandler(handler)
  
  
  def get_args():
      """gets command line arguments"""
  
      usage = """ Python script to resolve overlaps in ctms.  May be used with
                  utils/data/subsegment_data_dir.sh. """
      parser = argparse.ArgumentParser(usage)
      parser.add_argument('segments', type=argparse.FileType('r'),
                          help='use segments to resolve overlaps')
      parser.add_argument('ctm_in', type=argparse.FileType('r'),
                          help='input_ctm_file')
      parser.add_argument('ctm_out', type=argparse.FileType('w'),
                          help='output_ctm_file')
      parser.add_argument('--verbose', type=int, default=0,
                          help="Higher value for more verbose logging.")
      args = parser.parse_args()
  
      if args.verbose > 2:
          logger.setLevel(logging.DEBUG)
          handler.setLevel(logging.DEBUG)
  
      return args
  
  
  def read_segments(segments_file):
      """Read from segments and returns two dictionaries,
      {utterance-id: (recording_id, start_time, end_time)}
      {recording_id: list-of-utterances}
      """
      segments = {}
      reco2utt = defaultdict(list)
  
      num_lines = 0
      for line in segments_file:
          num_lines += 1
          parts = line.strip().split()
          assert len(parts) in [4, 5]
          segments[parts[0]] = (parts[1], float(parts[2]), float(parts[3]))
          reco2utt[parts[1]].append(parts[0])
  
      logger.info("Read %d lines from segments file %s",
                  num_lines, segments_file.name)
      segments_file.close()
  
      return segments, reco2utt
  
  
  def read_ctm(ctm_file, segments):
      """Read CTM from ctm_file into a dictionary of values indexed by the
      recording.
      It is assumed to be sorted by the recording-id and utterance-id.
  
      Returns a dictionary {recording : ctm_lines}
          where ctm_lines is a list of lines of CTM corresponding to the
          utterances in the recording.
          The format is as follows:
          [[(utteranceA, channelA, start_time1, duration1, hyp_word1, conf1),
            (utteranceA, channelA, start_time2, duration2, hyp_word2, conf2),
            ...
            (utteranceA, channelA, start_timeN, durationN, hyp_wordN, confN)],
           [(utteranceB, channelB, start_time1, duration1, hyp_word1, conf1),
            (utteranceB, channelB, start_time2, duration2, hyp_word2, conf2),
            ...],
           ...
           [...
            (utteranceZ, channelZ, start_timeN, durationN, hyp_wordN, confN)]
          ]
      """
      ctms = {}
  
      num_lines = 0
      for line in ctm_file:
          num_lines += 1
          parts = line.split()
  
          utt = parts[0]
          reco = segments[utt][0]
  
          if (reco, utt) not in ctms:
              ctms[(reco, utt)] = []
  
          ctms[(reco, utt)].append([parts[0], parts[1], float(parts[2]),
                                    float(parts[3])] + parts[4:])
  
      logger.info("Read %d lines from CTM %s", num_lines, ctm_file.name)
  
      ctm_file.close()
      return ctms
  
  
  def resolve_overlaps(ctms, segments):
      """Resolve overlaps within segments of the same recording.
  
      Returns new lines of CTM for the recording.
  
      Arguments:
          ctms - The CTM lines for a single recording. This is one value stored
              in the dictionary read by read_ctm(). Assumes that the lines
              are sorted by the utterance-ids.
              The format is the following:
              [[(utteranceA, channelA, start_time1, duration1, hyp_word1, conf1),
                (utteranceA, channelA, start_time2, duration2, hyp_word2, conf2),
                ...
                (utteranceA, channelA, start_timeN, durationN, hyp_wordN, confN)
               ],
               [(utteranceB, channelB, start_time1, duration1, hyp_word1, conf1),
                (utteranceB, channelB, start_time2, duration2, hyp_word2, conf2),
                ...],
               ...
               [...
                (utteranceZ, channelZ, start_timeN, durationN, hyp_wordN, confN)]
              ]
          segments - Dictionary containing the output of read_segments()
              { utterance_id: (recording_id, start_time, end_time) }
          """
      total_ctm = []
      if len(ctms) == 0:
          raise RuntimeError('CTMs for recording is empty. '
                             'Something wrong with the input ctms')
  
      # First column of first line in CTM for first utterance
      next_utt = ctms[0][0][0]
      for utt_index, ctm_for_cur_utt in enumerate(ctms):
          if utt_index == len(ctms) - 1:
              break
  
          if len(ctm_for_cur_utt) == 0:
              next_utt = ctms[utt_index + 1][0][0]
              continue
  
          cur_utt = ctm_for_cur_utt[0][0]
          if cur_utt != next_utt:
              logger.error(
                  "Current utterance %s is not the same as the next "
                  "utterance %s in previous iteration.
  "
                  "CTM is not sorted by utterance-id?",
                  cur_utt, next_utt)
              raise ValueError
  
          # Assumption here is that the segments are written in
          # consecutive order?
          ctm_for_next_utt = ctms[utt_index + 1]
          next_utt = ctm_for_next_utt[0][0]
          if segments[next_utt][1] < segments[cur_utt][1]:
              logger.error(
                  "Next utterance %s <= Current utterance %s. "
                  "CTM is not sorted by start-time of utterance-id.",
                  next_utt, cur_utt)
              raise ValueError
  
          try:
              # length of this utterance
              window_length = segments[cur_utt][2] - segments[cur_utt][1]
  
              # overlap of this segment with the next segment
              # i.e. current_utterance_end_time - next_utterance_start_time
              # Note: It is possible for this to be negative when there is
              # actually no overlap between consecutive segments.
              try:
                  overlap = segments[cur_utt][2] - segments[next_utt][1]
              except KeyError:
                  logger("Could not find utterance %s in segments",
                         next_utt)
                  raise
  
              if overlap > 0 and segments[next_utt][2] <= segments[cur_utt][2]:
                  # Next utterance is entirely within this utterance.
                  # So we leave this ctm as is and make the next one empty.
                  total_ctm.extend(ctm_for_cur_utt)
                  ctms[utt_index + 1] = []
                  continue
  
              # find a break point (a line in the CTM) for the current utterance
              # i.e. the first line that has more than half of it outside
              # the first half of the overlap region.
              # Note: This line will not be included in the output CTM, which is
              # only upto the line before this.
              try:
                  index = next(
                      (i for i, line in enumerate(ctm_for_cur_utt)
                       if (line[2] + line[3] / 2.0
                           > window_length - overlap / 2.0)))
              except StopIteration:
                  # It is possible for such a word to not exist, e.g the last
                  # word in the CTM is longer than overlap length and starts
                  # before the beginning of the overlap.
                  # or the last word ends before the middle of the overlap.
                  index = len(ctm_for_cur_utt)
  
              # Ignore the hypotheses beyond this midpoint. They will be
              # considered as part of the next segment.
              total_ctm.extend(ctm_for_cur_utt[:index])
  
              # Find a break point (a line in the CTM) for the next utterance
              # i.e. the first line that has more than half of it outside
              # the first half of the overlap region.
              try:
                  index = next(
                      (i for i, line in enumerate(ctm_for_next_utt)
                      if line[2] + line[3] / 2.0 > overlap / 2.0))
              except StopIteration:
                  # This can happen if there is no word hypothesized after
                  # half the overlap region.
                  ctms[utt_index + 1] = []
                  continue
  
              if index > 0:
                  # Update the ctm_for_next_utt to include only the lines
                  # starting from index.
                  ctms[utt_index + 1] = ctm_for_next_utt[index:]
              # else leave the ctm as is.
          except:
              logger.error("Could not resolve overlaps between CTMs for "
                           "%s and %s", cur_utt, next_utt)
              logger.error("Current CTM:")
              for line in ctm_for_cur_utt:
                  logger.error(ctm_line_to_string(line))
              logger.error("Next CTM:")
              for line in ctm_for_next_utt:
                  logger.error(ctm_line_to_string(line))
              raise
  
      # merge the last ctm entirely
      total_ctm.extend(ctms[-1])
  
      return total_ctm
  
  
  def ctm_line_to_string(line):
      """Converts a line of CTM to string."""
      return "{0} {1} {2} {3} {4}".format(line[0], line[1], line[2], line[3],
                                          " ".join(line[4:]))
  
  
  def write_ctm(ctm_lines, out_file):
      """Writes CTM lines stored in a list to file."""
      for line in ctm_lines:
          print(ctm_line_to_string(line), file=out_file)
  
  
  def run(args):
      """this method does everything in this script"""
      segments, reco2utt = read_segments(args.segments)
      ctms = read_ctm(args.ctm_in, segments)
  
      for reco, utts in reco2utt.items():
          ctms_for_reco = []
          for utt in sorted(utts, key=lambda x: segments[x][1]):
              if (reco, utt) in ctms:
                  ctms_for_reco.append(ctms[(reco, utt)])
          if len(ctms_for_reco) == 0:
              logger.info("CTM for recording {0} was empty".format(reco))
              continue
          try:
              # Process CTMs in the recordings
              ctms_for_reco = resolve_overlaps(ctms_for_reco, segments)
              write_ctm(ctms_for_reco, args.ctm_out)
          except Exception:
              logger.error("Failed to process CTM for recording %s",
                           reco)
              raise
      args.ctm_out.close()
      logger.info("Wrote CTM for %d recordings.", len(ctms))
  
  
  def main():
      """The main function which parses arguments and call run()."""
      args = get_args()
      try:
          run(args)
      except:
          logger.error("Failed to resolve overlaps", exc_info=True)
          raise SystemExit(1)
      finally:
          try:
              for f in [args.segments, args.ctm_in, args.ctm_out]:
                  if f is not None:
                      f.close()
          except IOError:
              logger.error("Could not close some files. "
                           "Disk error or broken pipes?")
              raise
          except UnboundLocalError:
              raise SystemExit(1)
  
  
  if __name__ == "__main__":
      main()