main.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. import os
  2. import warnings
  3. import tempfile
  4. from utils.files import filename, write_srt
  5. from utils.ffmpeg import get_audio, overlay_subtitles, add_subs_new
  6. from utils.bazarr import get_wanted_episodes, get_episode_details
  7. from utils.whisper import WhisperAI
  8. def process(args: dict):
  9. model_name: str = args.pop("model")
  10. output_dir: str = args.pop("output_dir")
  11. output_srt: bool = args.pop("output_srt")
  12. srt_only: bool = args.pop("srt_only")
  13. language: str = args.pop("language")
  14. sample_interval: str = args.pop("sample_interval")
  15. os.makedirs(output_dir, exist_ok=True)
  16. if model_name.endswith(".en"):
  17. warnings.warn(
  18. f"{model_name} is an English-only model, forcing English detection.")
  19. args["language"] = "en"
  20. # if translate task used and language argument is set, then use it
  21. elif language != "auto":
  22. args["language"] = language
  23. a = get_wanted_episodes()
  24. print(f"Found {a['total']} episodes needing subtitles.")
  25. for episode in a['data']:
  26. episode_data = get_episode_details(episode['sonarrEpisodeId'])
  27. print(episode_data)
  28. audios = get_audio(args.pop("video"), args.pop(
  29. 'audio_channel'), sample_interval)
  30. model_args = {}
  31. model_args["model_size_or_path"] = model_name
  32. model_args["device"] = args.pop("device")
  33. model_args["compute_type"] = args.pop("compute_type")
  34. srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir()
  35. subtitles = get_subtitles(audios, srt_output_dir, model_args, args)
  36. if srt_only:
  37. return
  38. add_subs_new(subtitles, output_dir, sample_interval)
  39. def get_subtitles(audio_paths: list, output_dir: str,
  40. model_args: dict, transcribe_args: dict):
  41. model = WhisperAI(model_args, transcribe_args)
  42. subtitles_path = {}
  43. for path, audio_path in audio_paths.items():
  44. print(
  45. f"Generating subtitles for {filename(path)}... This might take a while."
  46. )
  47. srt_path = os.path.join(output_dir, f"{filename(path)}.srt")
  48. segments = model.transcribe(audio_path)
  49. with open(srt_path, "w", encoding="utf-8") as srt:
  50. write_srt(segments, file=srt)
  51. subtitles_path[path] = srt_path
  52. return subtitles_path