deven367's picture
look for wav files
d1a4980
raw
history blame contribute delete
No virus
2.64 kB
from pathlib import Path
import numpy as np
import streamlit as st
from fastcore.xtras import globtastic
from annotator.utils import (
annotate,
df_from_result,
find_word_timestamp,
generate_srt,
get_audio,
get_v_from_url,
working_directory,
write_srt,
)
st.set_page_config(layout="wide")
SRT_PATH = Path("srt")
if not SRT_PATH.exists():
SRT_PATH.mkdir()
AUDIO_PATH = Path("./audio")
if not AUDIO_PATH.exists():
AUDIO_PATH.mkdir()
def make_sidebar():
with st.sidebar:
st.markdown("## yt-video-annotator")
st.write("Link to the GitHub repo")
def caption_from_url(url):
audio_src = get_audio(url)
v = get_v_from_url(url)
audio_src = globtastic(AUDIO_PATH, file_glob="*.wav", file_re=v)[0]
result = annotate(audio_src)
df = df_from_result(result)
return audio_src, df
def main():
url, name = None, None
make_sidebar()
place = "https://www.youtube.com/watch?v=C0DPdy98e4c&ab_channel=SimonYapp"
col1, col2 = st.columns([1.2, 1])
with col1:
url = st.text_input("Enter URL for the YT video", place)
st.video(url)
with col2:
default_opt = "Search for words"
opt = st.radio(
"What do you wish to do?",
[default_opt, "Generate subtitles for the entire video"],
)
if opt == default_opt:
st.markdown("### Search for words in the video")
words = st.text_input("Enter words separated by a comma")
words = words.split(",")
if st.button("Get Timestamps"):
audio_src, df = caption_from_url(url)
times = find_word_timestamp(df, *words)
times = np.asarray(times).reshape(len(words), -1)
# st.write(times)
for i, word in enumerate(words):
st.write(f"{word} is said on {times[i].flatten()} timestamp")
else:
if st.button("Generate SRT"):
audio_src, df = caption_from_url(url)
name = Path(audio_src).stem
s = generate_srt(df)
with working_directory(SRT_PATH):
write_srt(s, name)
if name is not None:
with working_directory(SRT_PATH):
key = get_v_from_url(url)
srt = globtastic(".", file_glob="*.srt", file_re=key)[0]
with open(srt) as f:
st.download_button("Download SRT", f, file_name=f"{name}.srt")
# subprocess.run(['rm', '-rf', 'audio'])
# subprocess.run(['rm', '-rf', 'srt'])
if __name__ == "__main__":
main()