BatuhanYilmaz commited on
Commit
f13254e
β€’
1 Parent(s): 0a8bf2e

Update 01_πŸŽ₯_Input_YouTube_Link.py

Browse files
Files changed (1) hide show
  1. 01_πŸŽ₯_Input_YouTube_Link.py +92 -92
01_πŸŽ₯_Input_YouTube_Link.py CHANGED
@@ -1,5 +1,6 @@
1
  import whisper
2
- from pytube import YouTube
 
3
  import requests
4
  import time
5
  import streamlit as st
@@ -11,17 +12,39 @@ from io import StringIO
11
  from utils import write_vtt, write_srt
12
  import ffmpeg
13
  from languages import LANGUAGES
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
 
15
- st.set_page_config(page_title="Auto Subtitled Video Generator ", page_icon=":movie_camera:", layout="wide")
16
 
17
  # Define a function that we can use to load lottie files from a link.
18
- @st.cache()
19
  def load_lottieurl(url: str):
20
  r = requests.get(url)
21
  if r.status_code != 200:
22
  return None
23
  return r.json()
24
 
 
 
 
 
 
 
 
 
 
25
  col1, col2 = st.columns([1, 3])
26
  with col1:
27
  lottie = load_lottieurl("https://assets8.lottiefiles.com/packages/lf20_jh9gfdye.json")
@@ -34,24 +57,12 @@ with col2:
34
  ###### ➠ If you want to transcribe the video in its original language, select the task as "Transcribe"
35
  ###### ➠ If you want to translate the subtitles to English, select the task as "Translate"
36
  ###### I recommend starting with the base model and then experimenting with the larger models, the small and medium models often work well. """)
37
-
38
 
39
- @st.cache(allow_output_mutation=True)
40
- def populate_metadata(link):
41
- yt = YouTube(link)
42
- author = yt.author
43
- title = yt.title
44
- description = yt.description
45
- thumbnail = yt.thumbnail_url
46
- length = yt.length
47
- views = yt.views
48
- return author, title, description, thumbnail, length, views
49
 
50
-
51
- @st.cache(allow_output_mutation=True)
52
  def download_video(link):
53
- yt = YouTube(link)
54
- video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first().download()
 
55
  return video
56
 
57
 
@@ -59,11 +70,6 @@ def convert(seconds):
59
  return time.strftime("%H:%M:%S", time.gmtime(seconds))
60
 
61
 
62
- loaded_model = whisper.load_model("base")
63
- current_size = "None"
64
-
65
-
66
- @st.cache(allow_output_mutation=True)
67
  def change_model(current_size, size):
68
  if current_size != size:
69
  loaded_model = whisper.load_model(size)
@@ -72,10 +78,10 @@ def change_model(current_size, size):
72
  raise Exception("Model size is the same as the current size.")
73
 
74
 
75
- @st.cache(allow_output_mutation=True)
76
  def inference(link, loaded_model, task):
77
- yt = YouTube(link)
78
- path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp3")
 
79
  if task == "Transcribe":
80
  options = dict(task="transcribe", best_of=5)
81
  results = loaded_model.transcribe(path, **options)
@@ -94,7 +100,6 @@ def inference(link, loaded_model, task):
94
  raise ValueError("Task not supported")
95
 
96
 
97
- @st.cache(allow_output_mutation=True)
98
  def getSubs(segments: Iterator[dict], format: str, maxLineWidth: int) -> str:
99
  segmentStream = StringIO()
100
 
@@ -120,8 +125,8 @@ def get_language_code(language):
120
  def generate_subtitled_video(video, audio, transcript):
121
  video_file = ffmpeg.input(video)
122
  audio_file = ffmpeg.input(audio)
123
- ffmpeg.concat(video_file.filter("subtitles", transcript), audio_file, v=1, a=1).output("final.mp4").run(quiet=True, overwrite_output=True)
124
- video_with_subs = open("final.mp4", "rb")
125
  return video_with_subs
126
 
127
 
@@ -130,25 +135,27 @@ def main():
130
  loaded_model = change_model(current_size, size)
131
  st.write(f"Model is {'multilingual' if loaded_model.is_multilingual else 'English-only'} "
132
  f"and has {sum(np.prod(p.shape) for p in loaded_model.parameters()):,} parameters.")
133
- link = st.text_input("YouTube Link (The longer the video, the longer the processing time)")
134
  task = st.selectbox("Select Task", ["Transcribe", "Translate"], index=0)
135
  if task == "Transcribe":
136
  if st.button("Transcribe"):
137
- author, title, description, thumbnail, length, views = populate_metadata(link)
138
- results = inference(link, loaded_model, task)
139
  video = download_video(link)
140
  lang = results[3]
141
  detected_language = get_language_code(lang)
142
 
143
  col3, col4 = st.columns(2)
144
- col5, col6, col7, col8 = st.columns(4)
145
- col9, col10 = st.columns(2)
146
  with col3:
147
  st.video(video)
148
-
149
- # Write the results to a .txt file and download it.
 
 
 
 
150
  with open("transcript.txt", "w+", encoding='utf8') as f:
151
- f.writelines(results[0])
152
  f.close()
153
  with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:
154
  datatxt = f.read()
@@ -164,50 +171,47 @@ def main():
164
  f.close()
165
  with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f:
166
  datasrt = f.read()
167
-
168
- with col5:
169
- st.download_button(label="Download Transcript (.txt)",
170
- data=datatxt,
171
- file_name="transcript.txt")
172
- with col6:
173
- st.download_button(label="Download Transcript (.vtt)",
174
- data=datavtt,
175
- file_name="transcript.vtt")
176
- with col7:
177
- st.download_button(label="Download Transcript (.srt)",
178
- data=datasrt,
179
- file_name="transcript.srt")
180
- with col9:
181
- st.success("You can download the transcript in .srt format, edit it (if you need to) and upload it to YouTube to create subtitles for your video.")
182
- with col10:
183
- st.info("Streamlit refreshes after the download button is clicked. The data is cached so you can download the transcript again without having to transcribe the video again.")
184
-
185
  with col4:
186
- with st.spinner("Generating Subtitled Video "):
187
- video_with_subs = generate_subtitled_video(video, "audio.mp3", "transcript.srt")
188
  st.video(video_with_subs)
189
  st.balloons()
190
- with col8:
191
- st.download_button(label="Download Subtitled Video",
192
- data=video_with_subs,
193
- file_name=f"{title} with subtitles.mp4")
 
 
 
 
 
 
 
 
 
 
 
 
194
  elif task == "Translate":
195
  if st.button("Translate to English"):
196
- author, title, description, thumbnail, length, views = populate_metadata(link)
197
- results = inference(link, loaded_model, task)
198
  video = download_video(link)
199
  lang = results[3]
200
  detected_language = get_language_code(lang)
201
 
202
  col3, col4 = st.columns(2)
203
- col5, col6, col7, col8 = st.columns(4)
204
- col9, col10 = st.columns(2)
205
  with col3:
206
  st.video(video)
207
 
208
- # Write the results to a .txt file and download it.
 
 
 
 
209
  with open("transcript.txt", "w+", encoding='utf8') as f:
210
- f.writelines(results[0])
211
  f.close()
212
  with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:
213
  datatxt = f.read()
@@ -223,36 +227,32 @@ def main():
223
  f.close()
224
  with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f:
225
  datasrt = f.read()
226
- with col5:
227
- st.download_button(label="Download Transcript (.txt)",
228
- data=datatxt,
229
- file_name="transcript.txt")
230
- with col6:
231
- st.download_button(label="Download Transcript (.vtt)",
232
- data=datavtt,
233
- file_name="transcript.vtt")
234
- with col7:
235
- st.download_button(label="Download Transcript (.srt)",
236
- data=datasrt,
237
- file_name="transcript.srt")
238
- with col9:
239
- st.success("You can download the transcript in .srt format, edit it (if you need to) and upload it to YouTube to create subtitles for your video.")
240
- with col10:
241
- st.info("Streamlit refreshes after the download button is clicked. The data is cached so you can download the transcript again without having to transcribe the video again.")
242
-
243
  with col4:
244
- with st.spinner("Generating Subtitled Video "):
245
- video_with_subs = generate_subtitled_video(video, "audio.mp3", "transcript.srt")
246
  st.video(video_with_subs)
247
  st.balloons()
248
- with col8:
249
- st.download_button(label="Download Subtitled Video ",
250
- data=video_with_subs,
251
- file_name=f"{title} with subtitles.mp4")
 
 
 
 
 
 
 
 
 
 
 
 
252
  else:
253
- st.error("Please select a task.")
254
 
255
 
256
  if __name__ == "__main__":
257
  main()
258
- st.markdown("###### Made with :heart: by [@BatuhanYΔ±lmaz](https://github.com/BatuhanYilmaz26) [![this is an image link](https://i.imgur.com/thJhzOO.png)](https://www.buymeacoffee.com/batuhanylmz)")
 
1
  import whisper
2
+ from pytubefix import YouTube
3
+ from pytubefix.cli import on_progress
4
  import requests
5
  import time
6
  import streamlit as st
 
12
  from utils import write_vtt, write_srt
13
  import ffmpeg
14
  from languages import LANGUAGES
15
+ import torch
16
+ from zipfile import ZipFile
17
+ from io import BytesIO
18
+ import base64
19
+ import pathlib
20
+ import re
21
+
22
+ st.set_page_config(page_title="Auto Subtitled Video Generator", page_icon=":movie_camera:", layout="wide")
23
+
24
+ torch.cuda.is_available()
25
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
26
+ # Model options: tiny, base, small, medium, large
27
+ loaded_model = whisper.load_model("small", device=DEVICE)
28
+ current_size = "None"
29
+
30
 
 
31
 
32
  # Define a function that we can use to load lottie files from a link.
 
33
  def load_lottieurl(url: str):
34
  r = requests.get(url)
35
  if r.status_code != 200:
36
  return None
37
  return r.json()
38
 
39
+ APP_DIR = pathlib.Path(__file__).parent.absolute()
40
+
41
+ LOCAL_DIR = APP_DIR / "local_youtube"
42
+ LOCAL_DIR.mkdir(exist_ok=True)
43
+ save_dir = LOCAL_DIR / "output"
44
+ save_dir.mkdir(exist_ok=True)
45
+
46
+
47
+
48
  col1, col2 = st.columns([1, 3])
49
  with col1:
50
  lottie = load_lottieurl("https://assets8.lottiefiles.com/packages/lf20_jh9gfdye.json")
 
57
  ###### ➠ If you want to transcribe the video in its original language, select the task as "Transcribe"
58
  ###### ➠ If you want to translate the subtitles to English, select the task as "Translate"
59
  ###### I recommend starting with the base model and then experimenting with the larger models, the small and medium models often work well. """)
 
60
 
 
 
 
 
 
 
 
 
 
 
61
 
 
 
62
  def download_video(link):
63
+ yt = YouTube(link, on_progress_callback=on_progress)
64
+ ys = yt.streams.get_highest_resolution()
65
+ video = ys.download(filename=f"{save_dir}/youtube_video.mp4")
66
  return video
67
 
68
 
 
70
  return time.strftime("%H:%M:%S", time.gmtime(seconds))
71
 
72
 
 
 
 
 
 
73
  def change_model(current_size, size):
74
  if current_size != size:
75
  loaded_model = whisper.load_model(size)
 
78
  raise Exception("Model size is the same as the current size.")
79
 
80
 
 
81
  def inference(link, loaded_model, task):
82
+ yt = YouTube(link, on_progress_callback=on_progress)
83
+ ys = yt.streams.get_audio_only()
84
+ path = ys.download(filename=f"{save_dir}/audio.mp3", mp3=True)
85
  if task == "Transcribe":
86
  options = dict(task="transcribe", best_of=5)
87
  results = loaded_model.transcribe(path, **options)
 
100
  raise ValueError("Task not supported")
101
 
102
 
 
103
  def getSubs(segments: Iterator[dict], format: str, maxLineWidth: int) -> str:
104
  segmentStream = StringIO()
105
 
 
125
  def generate_subtitled_video(video, audio, transcript):
126
  video_file = ffmpeg.input(video)
127
  audio_file = ffmpeg.input(audio)
128
+ ffmpeg.concat(video_file.filter("subtitles", transcript), audio_file, v=1, a=1).output("youtube_sub.mp4").run(quiet=True, overwrite_output=True)
129
+ video_with_subs = open("youtube_sub.mp4", "rb")
130
  return video_with_subs
131
 
132
 
 
135
  loaded_model = change_model(current_size, size)
136
  st.write(f"Model is {'multilingual' if loaded_model.is_multilingual else 'English-only'} "
137
  f"and has {sum(np.prod(p.shape) for p in loaded_model.parameters()):,} parameters.")
138
+ link = st.text_input("YouTube Link (The longer the video, the longer the processing time)", placeholder="Input YouTube link and press enter")
139
  task = st.selectbox("Select Task", ["Transcribe", "Translate"], index=0)
140
  if task == "Transcribe":
141
  if st.button("Transcribe"):
142
+ with st.spinner("Transcribing the video..."):
143
+ results = inference(link, loaded_model, task)
144
  video = download_video(link)
145
  lang = results[3]
146
  detected_language = get_language_code(lang)
147
 
148
  col3, col4 = st.columns(2)
 
 
149
  with col3:
150
  st.video(video)
151
+
152
+ # Split result["text"] on !,? and . , but save the punctuation
153
+ sentences = re.split("([!?.])", results[0])
154
+ # Join the punctuation back to the sentences
155
+ sentences = ["".join(i) for i in zip(sentences[0::2], sentences[1::2])]
156
+ text = "\n\n".join(sentences)
157
  with open("transcript.txt", "w+", encoding='utf8') as f:
158
+ f.writelines(text)
159
  f.close()
160
  with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:
161
  datatxt = f.read()
 
171
  f.close()
172
  with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f:
173
  datasrt = f.read()
174
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
175
  with col4:
176
+ with st.spinner("Generating Subtitled Video"):
177
+ video_with_subs = generate_subtitled_video(video, f"{save_dir}/audio.mp3", "transcript.srt")
178
  st.video(video_with_subs)
179
  st.balloons()
180
+
181
+ zipObj = ZipFile("YouTube_transcripts_and_video.zip", "w")
182
+ zipObj.write("transcript.txt")
183
+ zipObj.write("transcript.vtt")
184
+ zipObj.write("transcript.srt")
185
+ zipObj.write("youtube_sub.mp4")
186
+ zipObj.close()
187
+ ZipfileDotZip = "YouTube_transcripts_and_video.zip"
188
+ with open(ZipfileDotZip, "rb") as f:
189
+ datazip = f.read()
190
+ b64 = base64.b64encode(datazip).decode()
191
+ href = f"<a href=\"data:file/zip;base64,{b64}\" download='{ZipfileDotZip}'>\
192
+ Download Transcripts and Video\
193
+ </a>"
194
+ st.markdown(href, unsafe_allow_html=True)
195
+
196
  elif task == "Translate":
197
  if st.button("Translate to English"):
198
+ with st.spinner("Translating to English..."):
199
+ results = inference(link, loaded_model, task)
200
  video = download_video(link)
201
  lang = results[3]
202
  detected_language = get_language_code(lang)
203
 
204
  col3, col4 = st.columns(2)
 
 
205
  with col3:
206
  st.video(video)
207
 
208
+ # Split result["text"] on !,? and . , but save the punctuation
209
+ sentences = re.split("([!?.])", results[0])
210
+ # Join the punctuation back to the sentences
211
+ sentences = ["".join(i) for i in zip(sentences[0::2], sentences[1::2])]
212
+ text = "\n\n".join(sentences)
213
  with open("transcript.txt", "w+", encoding='utf8') as f:
214
+ f.writelines(text)
215
  f.close()
216
  with open(os.path.join(os.getcwd(), "transcript.txt"), "rb") as f:
217
  datatxt = f.read()
 
227
  f.close()
228
  with open(os.path.join(os.getcwd(), "transcript.srt"), "rb") as f:
229
  datasrt = f.read()
230
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
231
  with col4:
232
+ with st.spinner("Generating Subtitled Video"):
233
+ video_with_subs = generate_subtitled_video(video, f"{save_dir}/audio.mp3", "transcript.srt")
234
  st.video(video_with_subs)
235
  st.balloons()
236
+
237
+ zipObj = ZipFile("YouTube_transcripts_and_video.zip", "w")
238
+ zipObj.write("transcript.txt")
239
+ zipObj.write("transcript.vtt")
240
+ zipObj.write("transcript.srt")
241
+ zipObj.write("youtube_sub.mp4")
242
+ zipObj.close()
243
+ ZipfileDotZip = "YouTube_transcripts_and_video.zip"
244
+ with open(ZipfileDotZip, "rb") as f:
245
+ datazip = f.read()
246
+ b64 = base64.b64encode(datazip).decode()
247
+ href = f"<a href=\"data:file/zip;base64,{b64}\" download='{ZipfileDotZip}'>\
248
+ Download Transcripts and Video\
249
+ </a>"
250
+ st.markdown(href, unsafe_allow_html=True)
251
+
252
  else:
253
+ st.info("Please select a task.")
254
 
255
 
256
  if __name__ == "__main__":
257
  main()
258
+