transcribe.py (11780B)
1 #!/usr/bin/env python3 2 3 import argparse 4 import logging 5 import os 6 import shutil 7 import signal 8 import subprocess 9 from pathlib import Path 10 from typing import Any, Dict, List 11 12 import librosa 13 import matplotlib 14 import matplotlib.pyplot as plt 15 import tqdm 16 import yt_dlp 17 from ffmpeg_normalize import FFmpegNormalize 18 19 # Configure logging 20 logging.basicConfig( 21 level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" 22 ) 23 24 # Define constants 25 CONFIG_FILENAMES: List[str] = ["transcribe.cfg", "trancribe.cfg"] 26 DEFAULT_THREADS: int = 1 27 DEFAULT_MODEL: str = "medium" 28 DEFAULT_PLAYLIST_URL: str = "https://www.youtube.com/@Matrixdotorg" 29 SCRIPT_DIR: Path = Path(__file__).resolve().parent 30 31 32 # Define the function to handle Ctrl+C 33 def ctrl_c(sig: int, frame: Any) -> None: 34 logging.info("------------------------------") 35 logging.info("STOPPING...") 36 logging.info("bye bye") 37 logging.info("------------------------------") 38 exit(0) 39 40 41 # Set the signal handler 42 signal.signal(signal.SIGINT, ctrl_c) 43 44 45 def load_configuration(config_filenames: List[str]) -> Dict[str, str]: 46 config: Dict[str, str] = {} 47 for filename in config_filenames: 48 if Path(filename).is_file(): 49 with open(filename, "r") as cfg_file: 50 for line in cfg_file: 51 if line.strip() and not line.startswith("#"): 52 key, value = line.strip().split("=") 53 config[key.strip()] = value.strip() 54 return config 55 56 57 def create_config_file() -> None: 58 threads: int = get_thread_count() 59 model: str = ( 60 input("Enter the model (Press Enter for default): ").strip() or DEFAULT_MODEL 61 ) 62 63 # Save configuration 64 with open("transcribe.cfg", "w") as cfg_file: 65 cfg_file.write(f"THREADS={threads}\n") 66 cfg_file.write("PIDFILE=~/.trancribe-transcribe.pid\n") 67 cfg_file.write(f"MODEL={model}\n") 68 69 70 def check_dependencies() -> None: 71 dependencies: List[str] = ["ffmpeg", "git"] 72 for dependency in dependencies: 73 if shutil.which(dependency) is None: 74 logging.error(f"Can't find {dependency}. Please install.") 75 exit(1) 76 77 78 def setup_whisper() -> None: 79 whisper_dir: Path = SCRIPT_DIR / "whisper.cpp" 80 if not (whisper_dir / "whisper.cpp").exists(): 81 logging.info("Downloading whisper.cpp") 82 subprocess.run( 83 ["git", "clone", "https://github.com/ggerganov/whisper.cpp"], 84 cwd=SCRIPT_DIR, 85 check=True, 86 ) 87 logging.info("Compiling whisper...") 88 subprocess.run(["WHISPER_CLBLAST=1", "make", "-j"], cwd=whisper_dir, check=True) 89 logging.info("Downloading the model") 90 subprocess.run( 91 ["./models/download-ggml-model.sh", "medium"], cwd=whisper_dir, check=True 92 ) 93 else: 94 os.chdir(whisper_dir) 95 96 97 def get_thread_count() -> int: 98 threads: int = 0 99 try: 100 threads = os.cpu_count() or 0 101 except Exception as e: 102 logging.error(f"Error getting number of threads: {str(e)}") 103 104 if threads == 0: 105 logging.warn("Could not find the number of threads.") 106 else: 107 logging.info(f"Maximum of {threads} threads found.") 108 109 input_threads: str = input( 110 "How many threads to use for transcription? (Press Enter for default): " 111 ) 112 113 if input_threads.strip(): 114 try: 115 threads = int(input_threads) 116 except ValueError: 117 logging.error("Invalid input. Using default.") 118 return threads 119 120 121 def run_transcription_test(config: Dict[str, str]) -> None: 122 threads: int = int(config.get("THREADS", DEFAULT_THREADS)) 123 model: str = config.get("MODEL", DEFAULT_MODEL) 124 125 logging.info("Starting test. This might take some minutes, please wait...") 126 127 thread_opt: str = f"-t {threads}" if threads else "" 128 129 try: 130 subprocess.run( 131 [ 132 "./main", 133 "-m", 134 f"models/ggml-{model}.bin", 135 thread_opt, 136 "-l", 137 "de", 138 "-di", 139 "../test.wav", 140 ], 141 check=True, 142 ) 143 except subprocess.CalledProcessError: 144 logging.error("Error transcribing. Stopping.") 145 146 147 def download_audio_files(playlist_url: str, output_directory: str) -> None: 148 # Download audio files from YouTube playlist 149 ydl_opts: yt_dlp.YDLOpts = { 150 "format": "bestaudio/best", 151 "outtmpl": f"{output_directory}/%(id)s.%(ext)s", 152 "postprocessors": [{"key": "FFmpegExtractAudio", "preferredcodec": "wav"}], 153 "concurrent_fragment_downloads": 3, 154 "download_archive": f"{output_directory}/downloaded.txt", 155 "live_from_start": True, 156 "extractor_args": {"youtube": {"player_client": "android"}}, 157 } # type: ignore 158 159 with yt_dlp.YoutubeDL(ydl_opts) as ydl: 160 ydl.download([playlist_url]) 161 162 163 def normalize_audio_files(input_directory: str, output_directory: str) -> None: 164 # Normalize audio files using FFmpegNormalize 165 files: List[str] = [ 166 file for file in os.listdir(input_directory) if file.endswith(".wav") 167 ] 168 169 # Filter out files that are already in the output directory 170 existing_files = os.listdir(output_directory) 171 files_to_normalize = [file for file in files if file not in existing_files] 172 173 files_in: List[str] = [f"{input_directory}/{file}" for file in files_to_normalize] 174 files_out: List[str] = [f"{output_directory}/{file}" for file in files_to_normalize] 175 176 normalizer = FFmpegNormalize( 177 progress=True, 178 video_disable=True, 179 sample_rate=16000, 180 ) 181 for input_file, output_file in zip(files_in, files_out): 182 normalizer.add_media_file(input_file, output_file) 183 normalizer.run_normalization() 184 185 186 def generate_graph(audio_directory: str) -> None: 187 logging.info("Generating graph over audio files") 188 audio_files: List[str] = [ 189 file for file in os.listdir(audio_directory) if file.endswith(".wav") 190 ] 191 192 # Get lengths of audio files 193 audio_lengths: List[float] = [] 194 progress_bar = tqdm.tqdm(total=len(audio_files), desc="Processing audio files") 195 196 for file in audio_files: 197 audio_path = os.path.join(audio_directory, file) 198 duration = librosa.get_duration(path=audio_path) 199 audio_lengths.append(duration) 200 progress_bar.update(1) 201 202 progress_bar.close() 203 204 # Create histogram with specified bins 205 plt.figure(figsize=(15, 8)) 206 counts, bins, bars = plt.hist( 207 audio_lengths, 208 bins=50, 209 color="skyblue", 210 rwidth=0.7, 211 ) 212 213 plt.xlabel( 214 "Audio Length", 215 labelpad=15, 216 color="#333333", 217 fontname="Inter", 218 ) 219 plt.ylabel( 220 "Frequency", 221 labelpad=15, 222 color="#333333", 223 fontname="Inter", 224 ) 225 plt.title( 226 "Distribution of Audio Lengths", 227 pad=15, 228 color="#333333", 229 weight="bold", 230 fontname="Inter", 231 ) 232 233 # Format x-axis ticks for better readability 234 formatter = matplotlib.ticker.FuncFormatter( # type: ignore 235 lambda x, _: "{:.0f}h {:.0f}m".format(x // 3600, (x % 3600) // 60) 236 ) 237 plt.gca().xaxis.set_major_formatter(formatter) 238 plt.gca().set_xticks(bins) 239 plt.gca().set_xticklabels(plt.gca().get_xticklabels(), rotation=45, ha="right") 240 plt.gca().spines["top"].set_visible(False) 241 plt.gca().spines["right"].set_visible(False) 242 plt.gca().spines["left"].set_visible(False) 243 plt.gca().spines["bottom"].set_color("#DDDDDD") 244 plt.gca().tick_params(bottom=True, left=False) 245 plt.gca().set_axisbelow(True) 246 plt.gca().yaxis.grid(True, color="#EEEEEE") 247 plt.gca().xaxis.grid(False) 248 249 plt.tight_layout() 250 251 # Save the histogram as an image file 252 plt.savefig("../audio_lengths_histogram.png") 253 254 # Close the plot to release resources 255 plt.close() 256 257 258 def transcribe_audio_files(files_directory: str, model: str, threads: int) -> None: 259 files: List[str] = [ 260 file for file in os.listdir(files_directory) if file.endswith(".wav") 261 ] 262 263 for file in files: 264 input_file: str = f"{files_directory}/{file}" 265 base_filename: str = os.path.splitext(file)[0] 266 output_file: str = f"output/{base_filename}" 267 268 # Check if any of the output files already exist 269 txt_exists = os.path.exists(f"output/{base_filename}.txt") 270 srt_exists = os.path.exists(f"output/{base_filename}.srt") 271 vtt_exists = os.path.exists(f"output/{base_filename}.vtt") 272 if not (txt_exists and srt_exists and vtt_exists): 273 logging.info(f"Transcribing {file}") 274 whisper_cmd: List[str] = [ 275 "./main", 276 "-m", 277 f"models/ggml-{model}.bin", 278 "-t", 279 str(threads), 280 "-l", 281 "en", 282 "-otxt", 283 "-ovtt", 284 "-osrt", 285 "-pc", 286 "--file", 287 input_file, 288 "--output-file", 289 output_file, 290 "-et", 291 "3.0", 292 ] 293 294 whisper_process: subprocess.CompletedProcess = subprocess.run(whisper_cmd) 295 if whisper_process.returncode != 0: 296 logging.error("Error transcribing") 297 else: 298 logging.info("Transcription successful") 299 backup_file( 300 Path(input_file), 301 Path(input_file.replace("playlist_normalized", "playlist")), 302 Path("./backup"), 303 ) 304 305 306 def backup_file( 307 input_file: Path, normalized_file: Path, backup_directory: Path 308 ) -> None: 309 shutil.move(input_file, f"{backup_directory}/input/{input_file.name}") 310 logging.info(f"Moved input file {input_file} to backup directory") 311 312 shutil.move( 313 normalized_file, f"{backup_directory}/normalized/{normalized_file.name}" 314 ) 315 logging.info(f"Moved normalized file {normalized_file} to backup directory") 316 317 318 def run() -> None: 319 # Load configuration 320 config: Dict[str, str] = load_configuration(CONFIG_FILENAMES) 321 threads: int = int(config.get("THREADS", DEFAULT_THREADS)) 322 model: str = config.get("MODEL", DEFAULT_MODEL) 323 playlist_url: str = config.get("PLAYLIST_URL", DEFAULT_PLAYLIST_URL) 324 325 # Set LC_NUMERIC 326 os.environ["LC_NUMERIC"] = "en_US.UTF-8" 327 328 logging.info("Starting engines! Let's transcribe some episodes") 329 330 # Change directory to whisper.cpp 331 os.chdir("whisper.cpp") 332 333 # Create necessary directories if they don't exist 334 os.makedirs("./playlist", exist_ok=True) 335 open("./playlist/downloaded.txt", "a").close() 336 os.makedirs("./playlist_normalized", exist_ok=True) 337 os.makedirs("./backup", exist_ok=True) 338 os.makedirs("./backup/input", exist_ok=True) 339 os.makedirs("./backup/normalized", exist_ok=True) 340 341 download_audio_files(playlist_url, "./playlist") 342 343 generate_graph("./playlist") 344 345 # Normalize audio files 346 normalize_audio_files("./playlist", "./playlist_normalized") 347 348 # Transcribe audio files 349 transcribe_audio_files("./playlist_normalized", model, threads) 350 351 logging.info("Transcription completed!") 352 353 # Change back to the original directory 354 os.chdir("..") 355 356 357 def main() -> None: 358 parser: argparse.ArgumentParser = argparse.ArgumentParser( 359 description="Transcription script" 360 ) 361 parser.add_argument("--setup", help="Mode of operation") 362 args: argparse.Namespace = parser.parse_args() 363 364 if args.setup: 365 create_config_file() 366 config: Dict[str, str] = load_configuration(CONFIG_FILENAMES) 367 check_dependencies() 368 setup_whisper() 369 run_transcription_test(config) 370 else: 371 run() 372 373 374 if __name__ == "__main__": 375 main()