Exemple #1
0
def subset(input_path, output_path, audio_name, align_boundary_words, cer, wer, duration, gap, unk, num_speakers):
	cat = output_path.endswith('.json')
	meta = dict(
		align_boundary_words = align_boundary_words,
		cer = cer,
		wer = wer,
		duration = duration,
		gap = gap,
		unk = unk,
		num_speakers = num_speakers
	)
	transcript_cat = []
	for transcript_name in os.listdir(input_path):
		if not transcript_name.endswith('.json'):
			continue
		transcript = json.load(open(os.path.join(input_path, transcript_name)))
		transcript = [dict(meta = meta, **t) for t in transcripts.prune(transcript, audio_name = audio_name, **meta)]
		transcript_cat.extend(transcript)

		if not cat:
			os.makedirs(output_path, exist_ok = True)
			json.dump(
				transcript,
				open(os.path.join(output_path, transcript_name), 'w'),
				ensure_ascii = False,
				sort_keys = True,
				indent = 2
			)
	if cat:
		json.dump(transcript_cat, open(output_path, 'w'), ensure_ascii = False, sort_keys = True, indent = 2)
	print(output_path)
def main(args, ext_json=['.json', '.json.gz']):
    utils.enable_jit_fusion()

    assert args.output_json or args.output_html or args.output_txt or args.output_csv, \
     'at least one of the output formats must be provided'
    os.makedirs(args.output_path, exist_ok=True)

    audio_data_paths = set(
        p for f in args.input_path
        for p in ([os.path.join(f, g)
                   for g in os.listdir(f)] if os.path.isdir(f) else [f])
        if os.path.isfile(p) and any(map(p.endswith, args.ext)))
    json_data_paths = set(
        p for p in args.input_path if any(map(p.endswith, ext_json))
        and not utils.strip_suffixes(p, ext_json) in audio_data_paths)

    data_paths = list(audio_data_paths | json_data_paths)

    exclude = set([
        os.path.splitext(basename)[0]
        for basename in os.listdir(args.output_path)
        if basename.endswith('.json')
    ]) if args.skip_processed else None

    data_paths = [
        path for path in data_paths
        if exclude is None or os.path.basename(path) not in exclude
    ]

    text_pipeline, frontend, model, generator = setup(args)
    val_dataset = datasets.AudioTextDataset(
        data_paths, [text_pipeline],
        args.sample_rate,
        frontend=frontend if not args.frontend_in_model else None,
        mono=args.mono,
        time_padding_multiple=args.batch_time_padding_multiple,
        audio_backend=args.audio_backend,
        exclude=exclude,
        max_duration=args.transcribe_first_n_sec,
        mode='batched_channels'
        if args.join_transcript else 'batched_transcript',
        string_array_encoding=args.dataset_string_array_encoding,
        debug_short_long_records_features_from_whole_normalized_signal=args.
        debug_short_long_records_features_from_whole_normalized_signal)
    print('Examples count: ', len(val_dataset))
    val_meta = val_dataset.pop_meta()
    val_data_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=None,
        collate_fn=val_dataset.collate_fn,
        num_workers=args.num_workers)
    csv_sep = dict(tab='\t', comma=',')[args.csv_sep]
    csv_lines = []  # only used if args.output_csv is True

    oom_handler = utils.OomHandler(max_retries=args.oom_retries)
    for i, (meta, s, x, xlen, y, ylen) in enumerate(val_data_loader):
        print(f'Processing: {i}/{len(val_dataset)}')
        meta = [val_meta[t['example_id']] for t in meta]

        audio_path = meta[0]['audio_path']
        audio_name = transcripts.audio_name(audio_path)
        begin_end = [dict(begin=t['begin'], end=t['end']) for t in meta]
        begin = torch.tensor([t['begin'] for t in begin_end],
                             dtype=torch.float)
        end = torch.tensor([t['end'] for t in begin_end], dtype=torch.float)
        #TODO WARNING assumes frontend not in dataset
        if not args.frontend_in_model:
            print('\n' * 10 + 'WARNING\n' * 5)
            print(
                'transcribe.py assumes frontend in model, in other case time alignment was incorrect'
            )
            print('WARNING\n' * 5 + '\n')

        duration = x.shape[-1] / args.sample_rate
        channel = [t['channel'] for t in meta]
        speaker = [t['speaker'] for t in meta]
        speaker_name = [t['speaker_name'] for t in meta]

        if x.numel() == 0:
            print(f'Skipping empty [{audio_path}].')
            continue

        try:
            tic = time.time()
            y, ylen = y.to(args.device), ylen.to(args.device)
            log_probs, logits, olen = model(
                x.squeeze(1).to(args.device), xlen.to(args.device))

            print('Input:', audio_name)
            print('Input time steps:', log_probs.shape[-1],
                  '| target time steps:', y.shape[-1])
            print(
                'Time: audio {audio:.02f} sec | processing {processing:.02f} sec'
                .format(audio=sum(map(transcripts.compute_duration, meta)),
                        processing=time.time() - tic))

            ts: shaping.Bt = duration * torch.linspace(
                0, 1, steps=log_probs.shape[-1],
                device=log_probs.device).unsqueeze(0).expand(x.shape[0], -1)

            ref_segments = [[
                dict(channel=channel[i],
                     begin=begin_end[i]['begin'],
                     end=begin_end[i]['end'],
                     ref=text_pipeline.postprocess(
                         text_pipeline.preprocess(meta[i]['ref'])))
            ] for i in range(len(meta))]
            hyp_segments = [
                alternatives[0] for alternatives in generator.generate(
                    tokenizer=text_pipeline.tokenizer,
                    log_probs=log_probs,
                    begin=begin,
                    end=end,
                    output_lengths=olen,
                    time_stamps=ts,
                    segment_text_key='hyp',
                    segment_extra_info=[
                        dict(speaker=s, speaker_name=sn, channel=c)
                        for s, sn, c in zip(speaker, speaker_name, channel)
                    ])
            ]
            hyp_segments = [
                transcripts.map_text(text_pipeline.postprocess, hyp=hyp)
                for hyp in hyp_segments
            ]
            hyp, ref = '\n'.join(
                transcripts.join(hyp=h)
                for h in hyp_segments).strip(), '\n'.join(
                    transcripts.join(ref=r) for r in ref_segments).strip()
            if args.verbose:
                print('HYP:', hyp)
            print('CER: {cer:.02%}'.format(cer=metrics.cer(hyp=hyp, ref=ref)))

            tic_alignment = time.time()
            if args.align and y.numel() > 0:
                alignment: shaping.BY = ctc.alignment(
                    log_probs.permute(2, 0, 1),
                    y[:, 0, :],  # assumed that 0 channel is char labels
                    olen,
                    ylen[:, 0],
                    blank=text_pipeline.tokenizer.eps_id,
                    pack_backpointers=args.pack_backpointers)
                aligned_ts: shaping.Bt = ts.gather(1, alignment)

                ref_segments = [
                    alternatives[0] for alternatives in generator.generate(
                        tokenizer=text_pipeline.tokenizer,
                        log_probs=torch.nn.functional.one_hot(
                            y[:,
                              0, :], num_classes=log_probs.shape[1]).permute(
                                  0, 2, 1),
                        begin=begin,
                        end=end,
                        output_lengths=ylen,
                        time_stamps=aligned_ts,
                        segment_text_key='ref',
                        segment_extra_info=[
                            dict(speaker=s, speaker_name=sn, channel=c)
                            for s, sn, c in zip(speaker, speaker_name, channel)
                        ])
                ]
                ref_segments = [
                    transcripts.map_text(text_pipeline.postprocess, hyp=ref)
                    for ref in ref_segments
                ]
            oom_handler.reset()
        except:
            if oom_handler.try_recover(model.parameters()):
                print(f'Skipping {i} / {len(val_dataset)}')
                continue
            else:
                raise

        print('Alignment time: {:.02f} sec'.format(time.time() -
                                                   tic_alignment))

        ref_transcript, hyp_transcript = [
            sorted(transcripts.flatten(segments), key=transcripts.sort_key)
            for segments in [ref_segments, hyp_segments]
        ]

        if args.max_segment_duration:
            if ref:
                ref_segments = list(
                    transcripts.segment_by_time(ref_transcript,
                                                args.max_segment_duration))
                hyp_segments = list(
                    transcripts.segment_by_ref(hyp_transcript, ref_segments))
            else:
                hyp_segments = list(
                    transcripts.segment_by_time(hyp_transcript,
                                                args.max_segment_duration))
                ref_segments = [[] for _ in hyp_segments]

        #### HACK for diarization
        elif args.ref_transcript_path and args.join_transcript:
            audio_name_hack = audio_name.split('.')[0]
            #TODO: normalize ref field
            ref_segments = [[t] for t in sorted(transcripts.load(
                os.path.join(args.ref_transcript_path, audio_name_hack +
                             '.json')),
                                                key=transcripts.sort_key)]
            hyp_segments = list(
                transcripts.segment_by_ref(hyp_transcript,
                                           ref_segments,
                                           set_speaker=True,
                                           soft=False))
        #### END OF HACK

        has_ref = bool(transcripts.join(ref=transcripts.flatten(ref_segments)))

        transcript = []
        for hyp_transcript, ref_transcript in zip(hyp_segments, ref_segments):
            hyp, ref = transcripts.join(hyp=hyp_transcript), transcripts.join(
                ref=ref_transcript)

            transcript.append(
                dict(audio_path=audio_path,
                     ref=ref,
                     hyp=hyp,
                     speaker_name=transcripts.speaker_name(ref=ref_transcript,
                                                           hyp=hyp_transcript),
                     words=metrics.align_words(
                         *metrics.align_strings(hyp=hyp, ref=ref))
                     if args.align_words else [],
                     words_ref=ref_transcript,
                     words_hyp=hyp_transcript,
                     **transcripts.summary(hyp_transcript),
                     **(dict(cer=metrics.cer(hyp=hyp, ref=ref))
                        if has_ref else {})))

        transcripts.collect_speaker_names(transcript,
                                          set_speaker_data=True,
                                          num_speakers=2)

        filtered_transcript = list(
            transcripts.prune(transcript,
                              align_boundary_words=args.align_boundary_words,
                              cer=args.prune_cer,
                              duration=args.prune_duration,
                              gap=args.prune_gap,
                              allowed_unk_count=args.prune_unk,
                              num_speakers=args.prune_num_speakers))

        print('Filtered segments:', len(filtered_transcript), 'out of',
              len(transcript))

        if args.output_json:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.json')
            print(transcripts.save(transcript_path, filtered_transcript))

        if args.output_html:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.html')
            print(
                vis.transcript(transcript_path, args.sample_rate, args.mono,
                               transcript, filtered_transcript))

        if args.output_txt:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.txt')
            with open(transcript_path, 'w') as f:
                f.write(' '.join(t['hyp'].strip()
                                 for t in filtered_transcript))
            print(transcript_path)

        if args.output_csv:
            assert len({t['audio_path'] for t in filtered_transcript}) == 1
            audio_path = filtered_transcript[0]['audio_path']
            hyp = ' '.join(t['hyp'].strip() for t in filtered_transcript)
            begin = min(t['begin'] for t in filtered_transcript)
            end = max(t['end'] for t in filtered_transcript)
            csv_lines.append(
                csv_sep.join([audio_path, hyp,
                              str(begin),
                              str(end)]))

        if args.logits:
            logits_file_path = os.path.join(args.output_path,
                                            audio_name + '.pt')
            if args.logits_crop:
                begin_end = [
                    dict(
                        zip(['begin', 'end'], [
                            t['begin'] + c / float(o) * (t['end'] - t['begin'])
                            for c in args.logits_crop
                        ])) for o, t in zip(olen, begin_end)
                ]
                logits_crop = [slice(*args.logits_crop) for o in olen]
            else:
                logits_crop = [slice(int(o)) for o in olen]

            # TODO: filter ref / hyp by channel?
            torch.save([
                dict(audio_path=audio_path,
                     logits=l[..., logits_crop[i]],
                     **begin_end[i],
                     ref=ref,
                     hyp=hyp) for i, l in enumerate(logits.cpu())
            ], logits_file_path)
            print(logits_file_path)

        print('Done: {:.02f} sec\n'.format(time.time() - tic))

    if args.output_csv:
        transcript_path = os.path.join(args.output_path, 'transcripts.csv')
        with open(transcript_path, 'w') as f:
            f.write('\n'.join(csv_lines))
        print(transcript_path)
Exemple #3
0
    def __init__(
        self,
        data_paths: typing.List[str],
        text_pipelines: typing.List[text_processing.ProcessingPipeline],
        sample_rate: int,
        mode: str = DEFAULT_MODE,
        frontend: typing.Optional[torch.nn.Module] = None,
        speaker_names: typing.Optional[typing.List[str]] = None,
        max_audio_file_size: typing.Optional[float] = None,  #bytes
        min_duration: typing.Optional[float] = None,
        max_duration: typing.Optional[float] = None,
        max_num_channels: int = 2,
        mono: bool = True,
        audio_dtype: str = 'float32',
        time_padding_multiple: int = 1,
        audio_backend: typing.Optional[str] = None,
        exclude: typing.Optional[typing.Set] = None,
        bucket_fn: typing.Callable[[typing.List[typing.Dict]],
                                   int] = lambda transcript: 0,
        pop_meta: bool = False,
        string_array_encoding: str = 'utf_16_le',
        _print: typing.Callable = print,
        debug_short_long_records_features_from_whole_normalized_signal:
        bool = False,
        duration_from_transcripts: bool = False,
    ):
        self.debug_short_long_records_features_from_whole_normalized_signal = debug_short_long_records_features_from_whole_normalized_signal
        self.mode = mode
        self.min_duration = min_duration
        self.max_duration = max_duration
        self.max_audio_file_size = max_audio_file_size
        self.text_pipelines = text_pipelines
        self.frontend = frontend
        self.sample_rate = sample_rate
        self.time_padding_multiple = time_padding_multiple
        self.mono = mono
        self.audio_backend = audio_backend
        self.audio_dtype = audio_dtype

        data_paths = data_paths if isinstance(data_paths,
                                              list) else [data_paths]

        data_paths_ = []
        for data_path in data_paths:
            if os.path.isdir(data_path):
                data_paths_.extend(
                    os.path.join(data_path, filename) for filename in filter(
                        audio.is_audio, os.listdir(data_path)))
            else:
                data_paths_.append(data_path)
        data_paths = data_paths_

        tic = time.time()

        segments = []
        for path in data_paths:
            if audio.is_audio(path):
                assert self.mono or self.mode != AudioTextDataset.DEFAULT_MODE, 'Only mono audio files allowed as dataset input in default mode'
                if self.mono:
                    transcript = [
                        dict(audio_path=path,
                             channel=transcripts.channel_missing)
                    ]
                else:
                    transcript = [
                        dict(audio_path=path, channel=c)
                        for c in range(max_num_channels)
                    ]
            else:
                transcript = transcripts.load(path)
            segments.extend(transcript)

        _print('Dataset reading time: ', time.time() - tic)
        tic = time.time()

        # get_or_else required because dictionary could contain None values which we want to replace.
        # dict.get doesnt work in this case
        get_or_else = lambda dictionary, key, default: dictionary[
            key] if dictionary.get(key) is not None else default
        for t in segments:
            t['ref'] = get_or_else(t, 'ref', transcripts.ref_missing)
            t['begin'] = get_or_else(t, 'begin', transcripts.time_missing)
            t['end'] = get_or_else(t, 'end', transcripts.time_missing)
            t['channel'] = get_or_else(
                t, 'channel', transcripts.channel_missing
            ) if not self.mono else transcripts.channel_missing

        transcripts.collect_speaker_names(segments,
                                          speaker_names=speaker_names or [],
                                          num_speakers=max_num_channels,
                                          set_speaker_data=True)

        buckets = []
        grouped_segments = []
        transcripts_len = []
        speakers_len = []
        if self.mode == AudioTextDataset.DEFAULT_MODE:
            groupped_transcripts = ((i, [t]) for i, t in enumerate(segments))
        else:
            groupped_transcripts = itertools.groupby(
                sorted(segments, key=transcripts.group_key),
                transcripts.group_key)

        for group_key, transcript in groupped_transcripts:
            transcript = sorted(transcript, key=transcripts.sort_key)
            if self.mode == AudioTextDataset.BATCHED_CHANNELS_MODE:
                transcript = transcripts.join_transcript(
                    transcript,
                    self.mono,
                    duration_from_transcripts=duration_from_transcripts)

            if exclude is not None:
                allowed_audio_names = set(
                    transcripts.audio_name(t) for t in transcript
                    if transcripts.audio_name(t) not in exclude)
            else:
                allowed_audio_names = None

            transcript = transcripts.prune(
                transcript,
                allowed_audio_names=allowed_audio_names,
                duration=(
                    min_duration if min_duration is not None else 0.0,
                    max_duration if max_duration is not None else 24.0 * 3600,
                ),  #24h
                max_audio_file_size=max_audio_file_size)
            transcript = list(transcript)
            for t in transcript:
                t['example_id'] = AudioTextDataset.get_example_id(t)

            if len(transcript) == 0:
                continue

            bucket = bucket_fn(transcript)
            for t in transcript:
                t['bucket'] = bucket
                speakers_len.append(
                    len(t['speaker']) if (
                        isinstance(t['speaker'], list)) else 1)
            buckets.append(bucket)
            grouped_segments.extend(transcript)
            transcripts_len.append(len(transcript))

        _print('Dataset construction time: ', time.time() - tic)
        tic = time.time()

        self.bucket = torch.tensor(buckets, dtype=torch.short)
        self.audio_path = utils.TensorBackedStringArray(
            [t['audio_path'] for t in grouped_segments],
            encoding=string_array_encoding)
        self.ref = utils.TensorBackedStringArray(
            [t['ref'] for t in grouped_segments],
            encoding=string_array_encoding)
        self.begin = torch.tensor([t['begin'] for t in grouped_segments],
                                  dtype=torch.float64)
        self.end = torch.tensor([t['end'] for t in grouped_segments],
                                dtype=torch.float64)
        self.channel = torch.tensor([t['channel'] for t in grouped_segments],
                                    dtype=torch.int8)
        self.example_id = utils.TensorBackedStringArray(
            [t['example_id'] for t in grouped_segments],
            encoding=string_array_encoding)
        if self.mode == AudioTextDataset.BATCHED_CHANNELS_MODE:
            self.speaker = torch.tensor([
                speaker for t in grouped_segments for speaker in t['speaker']
            ],
                                        dtype=torch.int64)
        else:
            self.speaker = torch.tensor(
                [t['speaker'] for t in grouped_segments], dtype=torch.int64)
        self.speaker_len = torch.tensor(speakers_len, dtype=torch.int16)
        self.transcript_cumlen = torch.tensor(transcripts_len,
                                              dtype=torch.int16).cumsum(
                                                  dim=0, dtype=torch.int64)
        if pop_meta:
            self.meta = {}
        else:
            self.meta = {t['example_id']: t for t in grouped_segments}
        _print('Dataset tensors creation time: ', time.time() - tic)
Exemple #4
0
def main(args):
    utils.enable_jit_fusion()

    assert args.output_json or args.output_html or args.output_txt or args.output_csv, \
     'at least one of the output formats must be provided'
    os.makedirs(args.output_path, exist_ok=True)
    data_paths = [
        p for f in args.input_path
        for p in ([os.path.join(f, g)
                   for g in os.listdir(f)] if os.path.isdir(f) else [f])
        if os.path.isfile(p) and any(map(p.endswith, args.ext))
    ] + [
        p
        for p in args.input_path if any(map(p.endswith, ['.json', '.json.gz']))
    ]
    exclude = set([
        os.path.splitext(basename)[0]
        for basename in os.listdir(args.output_path)
        if basename.endswith('.json')
    ] if args.skip_processed else [])
    data_paths = [
        path for path in data_paths if os.path.basename(path) not in exclude
    ]

    labels, frontend, model, decoder = setup(args)
    val_dataset = datasets.AudioTextDataset(
        data_paths, [labels],
        args.sample_rate,
        frontend=None,
        segmented=True,
        mono=args.mono,
        time_padding_multiple=args.batch_time_padding_multiple,
        audio_backend=args.audio_backend,
        exclude=exclude,
        max_duration=args.transcribe_first_n_sec,
        join_transcript=args.join_transcript,
        string_array_encoding=args.dataset_string_array_encoding)
    num_examples = len(val_dataset)
    print('Examples count: ', num_examples)
    val_data_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=None,
        collate_fn=val_dataset.collate_fn,
        num_workers=args.num_workers)
    csv_sep = dict(tab='\t', comma=',')[args.csv_sep]
    output_lines = []  # only used if args.output_csv is True

    oom_handler = utils.OomHandler(max_retries=args.oom_retries)
    for i, (meta, s, x, xlen, y, ylen) in enumerate(val_data_loader):
        print(f'Processing: {i}/{num_examples}')

        meta = [val_dataset.meta.get(m['example_id']) for m in meta]
        audio_path = meta[0]['audio_path']

        if x.numel() == 0:
            print(f'Skipping empty [{audio_path}].')
            continue

        begin = meta[0]['begin']
        end = meta[0]['end']
        audio_name = transcripts.audio_name(audio_path)

        try:
            tic = time.time()
            y, ylen = y.to(args.device), ylen.to(args.device)
            log_probs, olen = model(
                x.squeeze(1).to(args.device), xlen.to(args.device))

            #speech = vad.detect_speech(x.squeeze(1), args.sample_rate, args.window_size, aggressiveness = args.vad, window_size_dilate = args.window_size_dilate)
            #speech = vad.upsample(speech, log_probs)
            #log_probs.masked_fill_(models.silence_space_mask(log_probs, speech, space_idx = labels.space_idx, blank_idx = labels.blank_idx), float('-inf'))

            decoded = decoder.decode(log_probs, olen)

            print('Input:', audio_name)
            print('Input time steps:', log_probs.shape[-1],
                  '| target time steps:', y.shape[-1])
            print(
                'Time: audio {audio:.02f} sec | processing {processing:.02f} sec'
                .format(audio=sum(
                    transcripts.compute_duration(t) for t in meta),
                        processing=time.time() - tic))

            ts = (x.shape[-1] / args.sample_rate) * torch.linspace(
                0, 1,
                steps=log_probs.shape[-1]).unsqueeze(0) + torch.FloatTensor(
                    [t['begin'] for t in meta]).unsqueeze(1)
            channel = [t['channel'] for t in meta]
            speaker = [t['speaker'] for t in meta]
            ref_segments = [[
                dict(channel=channel[i],
                     begin=meta[i]['begin'],
                     end=meta[i]['end'],
                     ref=labels.decode(y[i, 0, :ylen[i]].tolist()))
            ] for i in range(len(decoded))]
            hyp_segments = [
                labels.decode(decoded[i],
                              ts[i],
                              channel=channel[i],
                              replace_blank=True,
                              replace_blank_series=args.replace_blank_series,
                              replace_repeat=True,
                              replace_space=False,
                              speaker=speaker[i] if isinstance(
                                  speaker[i], str) else None)
                for i in range(len(decoded))
            ]

            ref, hyp = '\n'.join(
                transcripts.join(ref=r)
                for r in ref_segments).strip(), '\n'.join(
                    transcripts.join(hyp=h) for h in hyp_segments).strip()
            if args.verbose:
                print('HYP:', hyp)
            print('CER: {cer:.02%}'.format(cer=metrics.cer(hyp=hyp, ref=ref)))

            tic_alignment = time.time()
            if args.align and y.numel() > 0:
                #if ref_full:# and not ref:
                #	#assert len(set(t['channel'] for t in meta)) == 1 or all(t['type'] != 'channel' for t in meta)
                #	#TODO: add space at the end
                #	channel = torch.ByteTensor(channel).repeat_interleave(log_probs.shape[-1]).reshape(1, -1)
                #	ts = ts.reshape(1, -1)
                #	log_probs = log_probs.transpose(0, 1).unsqueeze(0).flatten(start_dim = -2)
                #	olen = torch.tensor([log_probs.shape[-1]], device = log_probs.device, dtype = torch.long)
                #	y = y_full[None, None, :].to(y.device)
                #	ylen = torch.tensor([[y.shape[-1]]], device = log_probs.device, dtype = torch.long)
                #	segments = [([], sum([h for r, h in segments], []))]

                alignment = ctc.alignment(
                    log_probs.permute(2, 0, 1),
                    y.squeeze(1),
                    olen,
                    ylen.squeeze(1),
                    blank=labels.blank_idx,
                    pack_backpointers=args.pack_backpointers)
                ref_segments = [
                    labels.decode(y[i, 0, :ylen[i]].tolist(),
                                  ts[i],
                                  alignment[i],
                                  channel=channel[i],
                                  speaker=speaker[i],
                                  key='ref',
                                  speakers=val_dataset.speakers)
                    for i in range(len(decoded))
                ]
            oom_handler.reset()
        except:
            if oom_handler.try_recover(model.parameters()):
                print(f'Skipping {i} / {num_examples}')
                continue
            else:
                raise

        print('Alignment time: {:.02f} sec'.format(time.time() -
                                                   tic_alignment))

        if args.max_segment_duration:
            ref_transcript, hyp_transcript = [
                list(sorted(sum(segments, []), key=transcripts.sort_key))
                for segments in [ref_segments, hyp_segments]
            ]
            if ref:
                ref_segments = list(
                    transcripts.segment(ref_transcript,
                                        args.max_segment_duration))
                hyp_segments = list(
                    transcripts.segment(hyp_transcript, ref_segments))
            else:
                hyp_segments = list(
                    transcripts.segment(hyp_transcript,
                                        args.max_segment_duration))
                ref_segments = [[] for _ in hyp_segments]

        transcript = [
            dict(audio_path=audio_path,
                 ref=ref,
                 hyp=hyp,
                 speaker=transcripts.speaker(ref=ref_transcript,
                                             hyp=hyp_transcript),
                 cer=metrics.cer(hyp=hyp, ref=ref),
                 words=metrics.align_words(hyp=hyp, ref=ref)[-1]
                 if args.align_words else [],
                 alignment=dict(ref=ref_transcript, hyp=hyp_transcript),
                 **transcripts.summary(hyp_transcript)) for ref_transcript,
            hyp_transcript in zip(ref_segments, hyp_segments)
            for ref, hyp in [(transcripts.join(ref=ref_transcript),
                              transcripts.join(hyp=hyp_transcript))]
        ]
        filtered_transcript = list(
            transcripts.prune(transcript,
                              align_boundary_words=args.align_boundary_words,
                              cer=args.cer,
                              duration=args.duration,
                              gap=args.gap,
                              unk=args.unk,
                              num_speakers=args.num_speakers))

        print('Filtered segments:', len(filtered_transcript), 'out of',
              len(transcript))

        if args.output_json:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.json')
            print(transcript_path)
            with open(transcript_path, 'w') as f:
                json.dump(filtered_transcript,
                          f,
                          ensure_ascii=False,
                          sort_keys=True,
                          indent=2)

        if args.output_html:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.html')
            print(transcript_path)
            vis.transcript(transcript_path, args.sample_rate, args.mono,
                           transcript, filtered_transcript)

        if args.output_txt:
            transcript_path = os.path.join(args.output_path,
                                           audio_name + '.txt')
            print(transcript_path)
            with open(transcript_path, 'w') as f:
                f.write(hyp)

        if args.output_csv:
            output_lines.append(
                csv_sep.join((audio_path, hyp, str(begin), str(end))) + '\n')

        print('Done: {:.02f} sec\n'.format(time.time() - tic))

    if args.output_csv:
        with open(os.path.join(args.output_path, 'transcripts.csv'), 'w') as f:
            f.writelines(output_lines)
Exemple #5
0
def errors(input_path,
           include=[],
           exclude=[],
           audio=False,
           output_path=None,
           sortdesc=None,
           topk=None,
           duration=None,
           cer=None,
           wer=None,
           mer=None,
           filter_transcripts=None,
           strip_audio_path_prefix=''):
    include, exclude = (set(
        sum([
            list(map(transcripts.audio_name, json.load(open(file_path))))
            if file_path.endswith('.json') else
            open(file_path).read().splitlines() for file_path in clude
        ], [])) for clude in [include, exclude])
    read_transcript = lambda path: list(
        filter(
            lambda r: (not include or r['audio_name'] in include) and
            (not exclude or r['audio_name'] not in exclude),
            json.load(open(path))
            if isinstance(path, str) else path)) if path is not None else []
    ours, theirs = list(
        transcripts.prune(read_transcript(input_path[0]),
                          duration=duration,
                          cer=cer,
                          wer=wer,
                          mer=mer)), [{
                              r['audio_name']: r
                              for r in read_transcript(transcript)
                          } for transcript in input_path[1:]]

    if filter_transcripts is None:
        if sortdesc is not None:
            filter_transcripts = lambda cat: list(
                sorted(cat, key=lambda utt: utt[0][sortdesc], reverse=True))
        else:
            filter_transcripts = lambda cat: cat

    cat = filter_transcripts([
        [a] + list(filter(None, [t.get(a['audio_name'], None)
                                 for t in theirs])) for a in ours
    ])[slice(topk)]
    cat_by_labels = collections.defaultdict(list)
    for c in cat:
        transcripts_by_labels = collections.defaultdict(list)
        for transcript in c:
            transcripts_by_labels[transcript['labels_name']] += c
        for labels_name, grouped_transcripts in transcripts_by_labels.items():
            cat_by_labels[labels_name] += grouped_transcripts

    # TODO: add sorting https://stackoverflow.com/questions/14267781/sorting-html-table-with-javascript
    html_path = output_path or (input_path[0] + '.html')

    f = open(html_path, 'w')
    f.write(
        '<html><meta charset="utf-8"><style> table{border-collapse:collapse; width: 100%;} audio {width:100%} .br{border-right:2px black solid} tr.first>td {border-top: 1px solid black} tr.any>td {border-top: 1px dashed black}  .nowrap{white-space:nowrap} th.col{width:80px}</style>'
    )
    f.write(
        '<body><table><tr><th></th><th class="col">cer_easy</th><th class="col">cer</th><th class="col">wer_easy</th><th class="col">wer</th><th class="col">mer</th><th></th></tr>'
    )
    f.write('<tr><td><strong>averages<strong></td></tr>')
    f.write('\n'.join(
        '<tr><td class="br">{input_name}</td><td>{cer_easy:.02%}</td><td>{cer:.02%}</td><td>{wer_easy:.02%}</td><td>{wer:.02%}</td><td>{mer:.02%}</td></tr>'
        .format(
            input_name=os.path.basename(input_path[i]),
            cer=metrics.nanmean(c, 'cer'),
            wer=metrics.nanmean(c, 'wer'),
            mer=metrics.nanmean(c, 'mer'),
            cer_easy=metrics.nanmean(c, 'words_easy_errors_easy.cer_pseudo'),
            wer_easy=metrics.nanmean(c, 'words_easy_errors_easy.wer_pseudo'),
        ) for i, c in enumerate(zip(*cat))))
    if len(cat_by_labels.keys()) > 1:
        for labels_name, labels_transcripts in cat_by_labels.items():
            f.write(
                f'<tr><td><strong>averages ({labels_name})<strong></td></tr>')
            f.write('\n'.join(
                '<tr><td class="br">{input_name}</td><td>{cer_easy:.02%}</td><td>{cer:.02%}</td><td>{wer_easy:.02%}</td><td>{wer:.02%}</td><td>{mer:.02%}</td></tr>'
                .format(
                    input_name=os.path.basename(input_path[i]),
                    cer=metrics.nanmean(c, 'cer'),
                    wer=metrics.nanmean(c, 'wer'),
                    mer=metrics.nanmean(c, 'mer_wordwise'),
                    cer_easy=metrics.nanmean(
                        c, 'words_easy_errors_easy.cer_pseudo'),
                    wer_easy=metrics.nanmean(
                        c, 'words_easy_errors_easy.wer_pseudo'),
                ) for i, c in enumerate(zip(*labels_transcripts))))
    f.write('<tr><td>&nbsp;</td></tr>')
    f.write('\n'.join(
        f'''<tr class="first"><td colspan="6">''' +
        (f'<audio controls src="{audio_data_uri(utt[0]["audio_path"][len(strip_audio_path_prefix):])}"></audio>'
         if audio else '') +
        f'<div class="nowrap">{utt[0]["audio_name"]}</div></td><td>{word_alignment(utt[0], ref = True, flat = True)}</td><td>{word_alignment(utt[0], ref = True, flat = True)}</td></tr>'
        + '\n'.join(
            '<tr class="any"><td class="br">{audio_name}</td><td>{cer_easy:.02%}</td><td>{cer:.02%}</td><td>{wer_easy:.02%}</td><td>{wer:.02%}</td><td class="br">{mer:.02%}</td><td>{word_alignment}</td><td>{word_alignment_flat}</td></tr>'
            .format(audio_name=transcripts.audio_name(input_path[i]),
                    cer_easy=a.get("words_easy_errors_easy", {}).get(
                        "cer_pseudo", -1),
                    cer=a.get("cer", 1),
                    wer_easy=a.get("words_easy_errors_easy", {}).get(
                        "wer_pseudo", -1),
                    wer=a.get("wer", 1),
                    mer=a.get("mer_wordwise", 1),
                    word_alignment=word_alignment(a.get('words', [])),
                    word_alignment_flat=word_alignment(a, hyp=True, flat=True))
            for i, a in enumerate(utt)) for utt in cat))
    f.write('</table></body></html>')
    print(html_path)