package io.lumify.sphinx;
import com.google.inject.Inject;
import io.lumify.core.ingest.graphProperty.GraphPropertyWorkData;
import io.lumify.core.ingest.graphProperty.GraphPropertyWorker;
import io.lumify.core.ingest.video.VideoTranscript;
import io.lumify.core.model.audit.AuditAction;
import io.lumify.core.model.properties.LumifyProperties;
import io.lumify.core.util.ProcessRunner;
import org.securegraph.Element;
import org.securegraph.Property;
import org.securegraph.Vertex;
import org.securegraph.mutation.ExistingElementMutation;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.util.Map;
import static com.google.common.base.Preconditions.checkNotNull;
public class SphinxGraphPropertyWorker extends GraphPropertyWorker {
private static final long BYTES_PER_SAMPLE = 2;
private static final long SAMPLES_PER_SECOND = 16000;
public static final String MULTI_VALUE_KEY = SphinxGraphPropertyWorker.class.getName();
private ProcessRunner processRunner;
@Override
public void execute(InputStream in, GraphPropertyWorkData data) throws Exception {
VideoTranscript transcript = extractTranscriptFromAudio(data.getLocalFile());
if (transcript == null) {
return;
}
ExistingElementMutation<Vertex> m = data.getElement().prepareMutation();
Map<String, Object> metadata = data.createPropertyMetadata();
metadata.put(LumifyProperties.META_DATA_TEXT_DESCRIPTION, "Audio Transcript");
addVideoTranscriptAsTextPropertiesToMutation(m, MULTI_VALUE_KEY, transcript, metadata, data.getVisibility());
Vertex v = m.save(getAuthorizations());
getAuditRepository().auditVertexElementMutation(AuditAction.UPDATE, m, v, MULTI_VALUE_KEY, getUser(), data.getVisibility());
getAuditRepository().auditAnalyzedBy(AuditAction.ANALYZED_BY, v, getClass().getSimpleName(), getUser(), v.getVisibility());
getGraph().flush();
pushVideoTranscriptTextPropertiesOnWorkQueue(data.getElement(), MULTI_VALUE_KEY, transcript);
}
@Override
public boolean isHandled(Element element, Property property) {
if (property == null) {
return false;
}
String mimeType = (String) property.getMetadata().get(LumifyProperties.MIME_TYPE.getPropertyName());
return !(mimeType == null || !mimeType.startsWith("audio"));
}
private VideoTranscript extractTranscriptFromAudio(File localFile) throws IOException, InterruptedException {
checkNotNull(localFile, "localFile cannot be null");
File wavFile = File.createTempFile("encode_wav_", ".wav");
File wavFileNoSilence = File.createTempFile("encode_wav_no_silence_", ".wav");
File wavFileNoHeaders = File.createTempFile("encode_wav_noheader_", ".wav");
try {
convertAudioTo16bit1Ch(localFile, wavFile);
removeSilenceFromBeginning(wavFile, wavFileNoSilence);
long silenceFileSizeDiff = wavFile.length() - wavFileNoSilence.length();
double timeOffsetInSec = (double) silenceFileSizeDiff / BYTES_PER_SAMPLE / SAMPLES_PER_SECOND;
WavFileUtil.fixWavHeaders(wavFileNoSilence, wavFileNoHeaders); // TODO patch sphinx to handle headers correctly
String sphinxOutput = runSphinx(wavFileNoHeaders);
return SphinxOutputParser.parse(sphinxOutput, timeOffsetInSec);
} finally {
wavFile.delete();
wavFileNoSilence.delete();
wavFileNoHeaders.delete();
}
}
private String runSphinx(File inFile) throws IOException, InterruptedException {
checkNotNull(inFile, "inFile cannot be null");
ByteArrayOutputStream out = new ByteArrayOutputStream();
try {
processRunner.execute(
"pocketsphinx_continuous",
new String[]{
"-infile", inFile.getAbsolutePath(),
"-time", "true"
},
out,
inFile.getAbsolutePath() + ": "
);
} finally {
out.close();
}
return new String(out.toByteArray());
}
private void removeSilenceFromBeginning(File inFile, File outFile) throws IOException, InterruptedException {
checkNotNull(inFile, "inFile cannot be null");
checkNotNull(outFile, "outFile cannot be null");
processRunner.execute(
"sox",
new String[]{
inFile.getAbsolutePath(),
outFile.getAbsolutePath(),
"silence", "1", "0.1", "1%", // remove silence from beginning. at least 0.1s of less than 1% volume
"pad", "1", "0" // pad 1 second of silence to beginning
},
null,
inFile.getAbsolutePath() + ": "
);
}
private void convertAudioTo16bit1Ch(File inputFile, File outputFile) throws IOException, InterruptedException {
checkNotNull(inputFile, "inputFile cannot be null");
checkNotNull(outputFile, "outputFile cannot be null");
processRunner.execute(
"ffmpeg",
new String[]{
"-y", // overwrite output files
"-i", inputFile.getAbsolutePath(),
"-acodec", "pcm_s16le",
"-ac", "1",
"-ar", Long.toString(SAMPLES_PER_SECOND),
outputFile.getAbsolutePath()
},
null,
inputFile.getAbsolutePath() + ": "
);
}
@Inject
public void setProcessRunner(ProcessRunner processRunner) {
this.processRunner = processRunner;
}
@Override
public boolean isLocalFileRequired() {
return true;
}
}