#!/usr/bin/env python

from __future__ import print_function

from kaldi.segmentation import NnetSAD, SegmentationProcessor
from kaldi.nnet3 import NnetSimpleComputationOptions
from kaldi.util.table import SequentialMatrixReader

# Construct SAD
model = NnetSAD.read_model("final.raw")
post = NnetSAD.read_average_posteriors("post_output.vec")
transform = NnetSAD.make_sad_transform(post)
graph = NnetSAD.make_sad_graph()
decodable_opts = NnetSimpleComputationOptions()
decodable_opts.extra_left_context = 79
decodable_opts.extra_right_context = 21
decodable_opts.extra_left_context_initial = 0
decodable_opts.extra_right_context_final = 0
decodable_opts.frames_per_chunk = 150
decodable_opts.acoustic_scale = 0.3
sad = NnetSAD(model, transform, graph, decodable_opts=decodable_opts)
seg = SegmentationProcessor(target_labels=[2])

# Define feature pipeline as a Kaldi rspecifier
feats_rspec = "ark:compute-mfcc-feats --config=mfcc.conf scp:wav.scp ark:- |"

# Segment
with SequentialMatrixReader(feats_rspec) as f, open ("segments", "w") as s:
    for key, feats in f:
        out = sad.segment(feats)
        segments, stats = seg.process(out["alignment"])
        seg.write(key, segments, s)
        print("segments:", segments, flush=True)
        print("stats:", stats, flush=True)
print("global stats:", seg.stats, flush=True)

[(10, 'kaldi.segmentation.NnetSAD.read_model', 'NnetSAD.read_model', (['"""final.raw"""'], {}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (11, 'kaldi.segmentation.NnetSAD.read_average_posteriors', 'NnetSAD.read_average_posteriors', (['"""post_output.vec"""'], {}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (12, 'kaldi.segmentation.NnetSAD.make_sad_transform', 'NnetSAD.make_sad_transform', (['post'], {}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (13, 'kaldi.segmentation.NnetSAD.make_sad_graph', 'NnetSAD.make_sad_graph', ([], {}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (14, 'kaldi.nnet3.NnetSimpleComputationOptions', 'NnetSimpleComputationOptions', ([], {}), False, 'from kaldi.nnet3 import NnetSimpleComputationOptions\n'), (21, 'kaldi.segmentation.NnetSAD', 'NnetSAD', (['model', 'transform', 'graph'], {'decodable_opts': 'decodable_opts'}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (22, 'kaldi.segmentation.SegmentationProcessor', 'SegmentationProcessor', ([], {'target_labels': '[2]'}), False, 'from kaldi.segmentation import NnetSAD, SegmentationProcessor\n'), (28, 'kaldi.util.table.SequentialMatrixReader', 'SequentialMatrixReader', (['feats_rspec'], {}), False, 'from kaldi.util.table import SequentialMatrixReader\n')]