|
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) |
|
|
|
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") |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
