# -*- coding: utf-8 -*-
from contextlib import contextmanager
import logging
import io
import subprocess
import sys
from datetime import timedelta
import ffmpeg
import numpy as np
from .sklearn_shim import TransformerMixin
from .sklearn_shim import Pipeline
import tqdm
from .constants import *
from .ffmpeg_utils import ffmpeg_bin_path, subprocess_args
from .subtitle_parser import make_subtitle_parser
from .subtitle_transformers import SubtitleScaler
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def make_subtitle_speech_pipeline(
fmt='srt',
encoding=DEFAULT_ENCODING,
caching=False,
max_subtitle_seconds=DEFAULT_MAX_SUBTITLE_SECONDS,
start_seconds=DEFAULT_START_SECONDS,
scale_factor=DEFAULT_SCALE_FACTOR,
parser=None,
**kwargs
):
if parser is None:
parser = make_subtitle_parser(
fmt,
encoding=encoding,
caching=caching,
max_subtitle_seconds=max_subtitle_seconds,
start_seconds=start_seconds
)
assert parser.encoding == encoding
assert parser.max_subtitle_seconds == max_subtitle_seconds
assert parser.start_seconds == start_seconds
return Pipeline([
('parse', parser),
('scale', SubtitleScaler(scale_factor)),
('speech_extract', SubtitleSpeechTransformer(
sample_rate=SAMPLE_RATE,
start_seconds=start_seconds,
framerate_ratio=scale_factor,
))
])
def _make_auditok_detector(sample_rate, frame_rate):
try:
from auditok import \
BufferAudioSource, ADSFactory, AudioEnergyValidator, StreamTokenizer
except ImportError as e:
logger.error("""Error: auditok not installed!
Consider installing it with `pip install auditok`. Note that auditok
is GPLv3 licensed, which means that successfully importing it at
runtime creates a derivative work that is GPLv3 licensed. For personal
use this is fine, but note that any commercial use that relies on
auditok must be open source as per the GPLv3!*
*Not legal advice. Consult with a lawyer.
""")
raise e
bytes_per_frame = 2
frames_per_window = frame_rate // sample_rate
validator = AudioEnergyValidator(
sample_width=bytes_per_frame, energy_threshold=50)
tokenizer = StreamTokenizer(
validator=validator, min_length=0.2*sample_rate,
max_length=int(5*sample_rate),
max_continuous_silence=0.25*sample_rate)
def _detect(asegment):
asource = BufferAudioSource(data_buffer=asegment,
sampling_rate=frame_rate,
sample_width=bytes_per_frame,
channels=1)
ads = ADSFactory.ads(audio_source=asource, block_dur=1./sample_rate)
ads.open()
tokens = tokenizer.tokenize(ads)
length = (len(asegment)//bytes_per_frame
+ frames_per_window - 1)//frames_per_window
media_bstring = np.zeros(length+1, dtype=int)
for token in tokens:
media_bstring[token[1]] += 1
media_bstring[token[2]+1] -= 1
return (np.cumsum(media_bstring)[:-1] > 0).astype(float)
return _detect
def _make_webrtcvad_detector(sample_rate, frame_rate):
import webrtcvad
vad = webrtcvad.Vad()
vad.set_mode(3) # set non-speech pruning aggressiveness from 0 to 3
window_duration = 1. / sample_rate # duration in seconds
frames_per_window = int(window_duration * frame_rate + 0.5)
bytes_per_frame = 2
def _detect(asegment):
media_bstring = []
failures = 0
for start in range(0, len(asegment) // bytes_per_frame,
frames_per_window):
stop = min(start + frames_per_window,
len(asegment) // bytes_per_frame)
try:
is_speech = vad.is_speech(
asegment[start * bytes_per_frame: stop * bytes_per_frame],
sample_rate=frame_rate)
except:
is_speech = False
failures += 1
# webrtcvad has low recall on mode 3, so treat non-speech as "not sure"
media_bstring.append(1. if is_speech else 0.5)
return np.array(media_bstring)
return _detect
class VideoSpeechTransformer(TransformerMixin):
def __init__(self, vad, sample_rate, frame_rate, start_seconds=0, ffmpeg_path=None, ref_stream=None, vlc_mode=False, gui_mode=False):
self.vad = vad
self.sample_rate = sample_rate
self.frame_rate = frame_rate
self.start_seconds = start_seconds
self.ffmpeg_path = ffmpeg_path
self.ref_stream = ref_stream
self.vlc_mode = vlc_mode
self.gui_mode = gui_mode
self.video_speech_results_ = None
def try_fit_using_embedded_subs(self, fname):
embedded_subs = []
embedded_subs_times = []
if self.ref_stream is None:
# check first 5; should cover 99% of movies
streams_to_try = map('0:s:{}'.format, range(5))
else:
streams_to_try = [self.ref_stream]
for stream in streams_to_try:
ffmpeg_args = [ffmpeg_bin_path('ffmpeg', self.gui_mode, ffmpeg_resources_path=self.ffmpeg_path)]
ffmpeg_args.extend([
'-loglevel', 'fatal',
'-nostdin',
'-i', fname,
'-map', '{}'.format(stream),
'-f', 'srt',
'-'
])
process = subprocess.Popen(ffmpeg_args, **subprocess_args(include_stdout=True))
output = io.BytesIO(process.communicate()[0])
if process.returncode != 0:
break
pipe = make_subtitle_speech_pipeline(start_seconds=self.start_seconds).fit(output)
speech_step = pipe.steps[-1][1]
embedded_subs.append(speech_step.subtitle_speech_results_)
embedded_subs_times.append(speech_step.max_time_)
if len(embedded_subs) == 0:
raise ValueError('Video file appears to lack subtitle stream')
# use longest set of embedded subs
self.video_speech_results_ = embedded_subs[int(np.argmax(embedded_subs_times))]
def fit(self, fname, *_):
if 'subs' in self.vad and (self.ref_stream is None or self.ref_stream.startswith('0:s:')):
try:
logger.info('Checking video for subtitles stream...')
self.try_fit_using_embedded_subs(fname)
logger.info('...success!')
return self
except Exception as e:
logger.info(e)
try:
total_duration = float(ffmpeg.probe(
fname, cmd=ffmpeg_bin_path('ffprobe', self.gui_mode, ffmpeg_resources_path=self.ffmpeg_path)
)['format']['duration']) - self.start_seconds
except Exception as e:
logger.warning(e)
total_duration = None
if 'webrtc' in self.vad:
detector = _make_webrtcvad_detector(self.sample_rate, self.frame_rate)
elif 'auditok' in self.vad:
detector = _make_auditok_detector(self.sample_rate, self.frame_rate)
else:
raise ValueError('unknown vad: %s' % self.vad)
media_bstring = []
ffmpeg_args = [ffmpeg_bin_path('ffmpeg', self.gui_mode, ffmpeg_resources_path=self.ffmpeg_path)]
if self.start_seconds > 0:
ffmpeg_args.extend([
'-ss', str(timedelta(seconds=self.start_seconds)),
])
ffmpeg_args.extend([
'-loglevel', 'fatal',
'-nostdin',
'-i', fname
])
if self.ref_stream is not None and self.ref_stream.startswith('0:a:'):
ffmpeg_args.extend(['-map', self.ref_stream])
ffmpeg_args.extend([
'-f', 's16le',
'-ac', '1',
'-acodec', 'pcm_s16le',
'-ar', str(self.frame_rate),
'-'
])
process = subprocess.Popen(ffmpeg_args, **subprocess_args(include_stdout=True))
bytes_per_frame = 2
frames_per_window = bytes_per_frame * self.frame_rate // self.sample_rate
windows_per_buffer = 10000
simple_progress = 0.
@contextmanager
def redirect_stderr(enter_result=None):
yield enter_result
tqdm_extra_args = {}
should_print_redirected_stderr = self.gui_mode
if self.gui_mode:
try:
from contextlib import redirect_stderr
tqdm_extra_args['file'] = sys.stdout
except ImportError:
should_print_redirected_stderr = False
pbar_output = io.StringIO()
with redirect_stderr(pbar_output):
with tqdm.tqdm(total=total_duration, disable=self.vlc_mode, **tqdm_extra_args) as pbar:
while True:
in_bytes = process.stdout.read(frames_per_window * windows_per_buffer)
if not in_bytes:
break
newstuff = len(in_bytes) / float(bytes_per_frame) / self.frame_rate
if total_duration is not None and simple_progress + newstuff > total_duration:
newstuff = total_duration - simple_progress
simple_progress += newstuff
pbar.update(newstuff)
if self.vlc_mode and total_duration is not None:
print("%d" % int(simple_progress * 100. / total_duration))
sys.stdout.flush()
if should_print_redirected_stderr:
assert self.gui_mode
# no need to flush since we pass -u to do unbuffered output for gui mode
print(pbar_output.read())
in_bytes = np.frombuffer(in_bytes, np.uint8)
media_bstring.append(detector(in_bytes))
if len(media_bstring) == 0:
raise ValueError(
'Unable to detect speech. Perhaps try specifying a different stream / track, or a different vad.'
)
self.video_speech_results_ = np.concatenate(media_bstring)
return self
def transform(self, *_):
return self.video_speech_results_
class SubtitleSpeechTransformer(TransformerMixin):
def __init__(self, sample_rate, start_seconds=0, framerate_ratio=1.):
self.sample_rate = sample_rate
self.start_seconds = start_seconds
self.framerate_ratio = framerate_ratio
self.subtitle_speech_results_ = None
self.max_time_ = None
def fit(self, subs, *_):
max_time = 0
for sub in subs:
max_time = max(max_time, sub.end.total_seconds())
self.max_time_ = max_time - self.start_seconds
samples = np.zeros(int(max_time * self.sample_rate) + 2, dtype=float)
for sub in subs:
start = int(round((sub.start.total_seconds() - self.start_seconds) * self.sample_rate))
duration = sub.end.total_seconds() - sub.start.total_seconds()
end = start + int(round(duration * self.sample_rate))
samples[start:end] = min(1. / self.framerate_ratio, 1.)
self.subtitle_speech_results_ = samples
return self
def transform(self, *_):
return self.subtitle_speech_results_
class DeserializeSpeechTransformer(TransformerMixin):
def __init__(self):
self.deserialized_speech_results_ = None
def fit(self, fname, *_):
speech = np.load(fname)
if hasattr(speech, 'files'):
if 'speech' in speech.files:
speech = speech['speech']
else:
raise ValueError('could not find "speech" array in '
'serialized file; only contains: %s' % speech.files)
self.deserialized_speech_results_ = speech
return self
def transform(self, *_):
return self.deserialized_speech_results_