dasemu/Transcriptarr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update subgen.py

Browse Source
This commit is contained in:
McCloudS
2023-10-22 14:18:00 -06:00
committed by GitHub
parent 33f40d464b
commit 8c1647413c
1 changed files with 46 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
subgen/subgen.py
View File

@@ -7,28 +7,35 @@ import threading
import av import av
import sys import sys
import time import time
import queue
def convert_to_bool(in_bool): def convert_to_bool(in_bool):
value = str(in_bool).lower() if isinstance(in_bool, bool):
return value not in ('false', 'off', '0') return in_bool
else:
value = str(in_bool).lower()
return value not in ('false', 'off', '0')
# Replace your getenv calls with appropriate default values here # Replace your getenv calls with appropriate default values here
plextoken = os.getenv('PLEXTOKEN', "token here") plextoken = os.getenv('PLEXTOKEN', "token here")
plexserver = os.getenv('PLEXSERVER', "http://192.168.1.111:32400") plexserver = os.getenv('PLEXSERVER', "http://192.168.1.111:32400")
whisper_model = os.getenv('WHISPER_MODEL', "medium") whisper_model = os.getenv('WHISPER_MODEL', "medium")
whisper_threads = int(os.getenv('WHISPER_THREADS', 4)) whisper_threads = int(os.getenv('WHISPER_THREADS', 4))
concurrent_transcriptions = int(os.getenv('CONCURRENT_TRANSCRIPTIONS', '1')) concurrent_transcriptions = int(os.getenv('CONCURRENT_TRANSCRIPTIONS', '2'))
procaddedmedia = convert_to_bool(os.getenv('PROCADDEDMEDIA', "True")) transcribe_device = os.getenv('TRANSCRIBE_DEVICE', "cpu")
procmediaonplay = convert_to_bool(os.getenv('PROCMEDIAONPLAY', "True")) procaddedmedia = convert_to_bool(os.getenv('PROCADDEDMEDIA', True))
procmediaonplay = convert_to_bool(os.getenv('PROCMEDIAONPLAY', True))
namesublang = os.getenv('NAMESUBLANG', "aa") namesublang = os.getenv('NAMESUBLANG', "aa")
skipifinternalsublang = os.getenv('SKIPIFINTERNALSUBLANG', "eng") skipifinternalsublang = os.getenv('SKIPIFINTERNALSUBLANG', "eng")
webhookport = int(os.getenv('WEBHOOKPORT', 8090)) webhookport = int(os.getenv('WEBHOOKPORT', 8090))
word_level_highlight = convert_to_bool(os.getenv('WORD_LEVEL_HIGHLIGHT', "False")) word_level_highlight = convert_to_bool(os.getenv('WORD_LEVEL_HIGHLIGHT', False))
debug = convert_to_bool(os.getenv('DEBUG', False)) debug = convert_to_bool(os.getenv('DEBUG', False))
use_path_mapping = convert_to_bool(os.getenv('USE_PATH_MAPPING', False)) use_path_mapping = convert_to_bool(os.getenv('USE_PATH_MAPPING', False))
path_mapping_from = os.getenv('PATH_MAPPING_FROM', '/tv') path_mapping_from = os.getenv('PATH_MAPPING_FROM', '/tv')
path_mapping_to = os.getenv('PATH_MAPPING_TO', '/Volumes/TV') path_mapping_to = os.getenv('PATH_MAPPING_TO', '/Volumes/TV')
store_local = convert_to_bool(os.getenv('STORE_LOCAL_LIBS', True)) store_local = convert_to_bool(os.getenv('STORE_LOCAL_LIBS', True))
if transcribe_device == "gpu":
transcribe_device = "cuda"
def install_if_not_installed(package_name): def install_if_not_installed(package_name):
try: try:
@@ -61,9 +68,11 @@ import stable_whisper
import requests import requests
app = Flask(__name__) app = Flask(__name__)
model = stable_whisper.load_faster_whisper(whisper_model, cpu_threads=whisper_threads) model = stable_whisper.load_faster_whisper(whisper_model, device=transcribe_device, cpu_threads=whisper_threads, num_workers=concurrent_transcriptions)
files_to_transcribe = set() files_to_transcribe = set()
subextension = '.subgen.' + whisper_model + '.' + namesublang + '.srt' subextension = '.subgen.' + whisper_model + '.' + namesublang + '.srt'
print("Transcriptions are limited to running " + str(concurrent_transcriptions) + " at a time")
print("Running " + str(whisper_threads) + " threads per transcription")
@app.route("/webhook", methods=["POST"]) @app.route("/webhook", methods=["POST"])
def receive_webhook(): def receive_webhook():
@@ -89,30 +98,33 @@ def receive_webhook():
print("Didn't get a webhook we expected, discarding") print("Didn't get a webhook we expected, discarding")
return "" return ""
print("Path of file: " + fullpath) if debug:
print("Path of file: " + fullpath)
print("event: " + event)
if use_path_mapping: if use_path_mapping:
fullpath = fullpath.replace(path_mapping_from, path_mapping_to) fullpath = fullpath.replace(path_mapping_from, path_mapping_to)
print("Updated path: " + fullpath.replace(path_mapping_from, path_mapping_to)) if debug:
print("event: " + event) print("Updated path: " + fullpath.replace(path_mapping_from, path_mapping_to))
print("Transcriptions are limited to running " + str(concurrent_transcriptions) + " at a time")
print("Running " + str(whisper_threads) + " threads per transcription")
add_file_for_transcription(fullpath) add_file_for_transcription(fullpath)
return "" return ""
def gen_subtitles(inputvideo): def gen_subtitles(video_file_path: str) -> None:
try: try:
print(f"Transcribing file: {inputvideo}") print(f"Transcribing file: {video_file_path}")
start_time = time.time() start_time = time.time()
result = model.transcribe_stable(inputvideo) result = model.transcribe_stable(video_file_path)
result.to_srt_vtt(inputvideo.rsplit('.', 1)[0] + subextension, word_level=word_level_highlight) result.to_srt_vtt(video_file_path.rsplit('.', 1)[0] + subextension, word_level=word_level_highlight)
elapsed_time = time.time() - start_time elapsed_time = time.time() - start_time
minutes, seconds = divmod(int(elapsed_time), 60) minutes, seconds = divmod(int(elapsed_time), 60)
print(f"Transcription of {file_path} is completed, it took {minutes} minutes and {seconds} seconds to complete.") print(f"Transcription of {file_path} is completed, it took {minutes} minutes and {seconds} seconds to complete.")
files_to_transcribe.remove(inputvideo) files_to_transcribe.remove(video_file_path)
except Exception as e: except Exception as e:
print(f"Error processing or transcribing {file_path}: {e}") print(f"Error processing or transcribing {video_file_path}: {e}")
finally:
files_to_transcribe.remove(video_file_path)
# Function to add a file for transcription # Function to add a file for transcription
def add_file_for_transcription(file_path): def add_file_for_transcription(file_path):
@@ -129,8 +141,8 @@ def add_file_for_transcription(file_path):
print(f"Added {file_path} for transcription.") print(f"Added {file_path} for transcription.")
# Start transcription for the file in a separate thread # Start transcription for the file in a separate thread
transcription_thread = threading.Thread(target=gen_subtitles, args=(file_path,)) gen_subtitles(file_path)
transcription_thread.start()
else: else:
print(f"File {file_path} is already in the transcription list. Skipping.") print(f"File {file_path} is already in the transcription list. Skipping.")
@@ -158,19 +170,29 @@ def has_subtitle_language(video_file, target_language):
print(f"An error occurred: {e}") print(f"An error occurred: {e}")
return False return False
def get_file_name(item_id, plexserver, plextoken): def get_file_name(item_id: str, plexserver: str, plextoken: str) -> str:
"""Gets the full path to a file from the Plex server.
Args:
item_id: The ID of the item in the Plex library.
plexserver: The URL of the Plex server.
plextoken: The Plex token.
Returns:
The full path to the file.
"""
url = f"{plexserver}/library/metadata/{item_id}" url = f"{plexserver}/library/metadata/{item_id}"
response = requests.get(url, headers={ response = requests.get(url, headers={
"X-Plex-Token": "{plextoken}"}) "X-Plex-Token": "{plextoken}"})
if response.status_code == 200: if response.status_code == 200:
root = ET.fromstring(response.text) root = ET.fromstring(response.text)
fullpath = root.find(".//Part").attrib['file'] fullpath = root.find(".//Part").attrib['file']
return fullpath return fullpath
else: else:
print(f"Error: {response.text}") return ""
return
print("Starting webhook!") print("Starting webhook!")
if __name__ == "__main__": if __name__ == "__main__":