Advertisement
jalarab

ImportTV

Oct 31st, 2024
111
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 7.56 KB | None | 0 0
  1. import os
  2. import subprocess
  3. import sys
  4. import json
  5. import argparse
  6. import openai
  7. from tqdm import tqdm
  8.  
  9. client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
  10.  
  11. def should_transcode(input_file):
  12.     command = [
  13.         "ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=bit_rate,height", "-of", "json", input_file
  14.     ]
  15.     try:
  16.         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, check=True)
  17.         video_info = json.loads(result.stdout)
  18.         video_stream = video_info['streams'][0]
  19.         video_bitrate = int(video_stream.get('bit_rate', 0))
  20.         video_height = int(video_stream.get('height', 0))
  21.  
  22.         command_audio = [
  23.             "ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries", "stream=bit_rate", "-of", "json", input_file
  24.         ]
  25.         result_audio = subprocess.run(command_audio, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, check=True)
  26.         audio_info = json.loads(result_audio.stdout)
  27.         audio_stream = audio_info['streams'][0]
  28.         audio_bitrate = int(audio_stream.get('bit_rate', 0))
  29.  
  30.         should_transcode = video_bitrate > 1000000 or video_height > 540 or audio_bitrate > 128000
  31.         print(f"Debug: Video bitrate: {video_bitrate}, Video height: {video_height}, Audio bitrate: {audio_bitrate}, Should transcode: {should_transcode}", flush=True)
  32.         return should_transcode
  33.     except subprocess.CalledProcessError as e:
  34.         print(f"Error al verificar el archivo {input_file}: {e.stderr}")
  35.         return False
  36.  
  37. def transcode_video(input_file, temp_file):
  38.     print("Debug: Iniciando transcodificación...", flush=True)
  39.     print(f"Iniciando transcodificación del archivo: {input_file}", flush=True)
  40.     command = [
  41.         "ffmpeg", "-i", input_file,
  42.         "-vf", "scale=-1:540,fps=fps=25",
  43.         "-c:v", "h264_qsv", "-preset", "veryfast", "-b:v", "1000k", "-minrate", "1000k", "-maxrate", "1000k", "-bufsize", "2800k", "-x264-params", "nal-hrd=cbr",
  44.         "-c:a", "aac", "-b:a", "128k", "-ac", "2",
  45.         "-map", "0",
  46.         "-c:s", "copy",
  47.         "-f", "matroska",
  48.         temp_file
  49.     ]
  50.  
  51.     try:
  52.         print(f"Debug: Ejecutando comando: {' '.join(command)}", flush=True)
  53.         process = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, check=True)
  54.         print(f"FFmpeg stderr output: {process.stderr}", flush=True)
  55.     except subprocess.CalledProcessError as e:
  56.         print(f"Error al transcodificar el archivo {input_file}: {e.stderr}")
  57.         return False
  58.     return True
  59.  
  60. def translate_text(text, target_language="es"):
  61.     chat_completion = client.chat.completions.create(
  62.         messages=[
  63.             {
  64.                 "role": "system",
  65.                 "content": f"Translate the following text to {target_language}.",
  66.             },
  67.             {
  68.                 "role": "user",
  69.                 "content": text,
  70.             },
  71.         ],
  72.         model="gpt-4o-mini",
  73.     )
  74.     translation = chat_completion.choices[0].message.content.strip()
  75.     return translation
  76.  
  77. def find_subtitle_file(input_file):
  78.     base_name = os.path.splitext(input_file)[0]
  79.     for ext in [".srt", ".en.srt", ".eng.srt", ".es.srt", ".spa.srt"]:
  80.         potential_file = f"{base_name}{ext}"
  81.         if os.path.exists(potential_file):
  82.             return potential_file
  83.     return None
  84.  
  85. def has_embedded_subtitles(input_file, language_code="es"):
  86.     command = [
  87.         "ffprobe", "-v", "error", "-select_streams", "s", "-show_entries", "stream_tags=language", "-of", "json", input_file
  88.     ]
  89.     try:
  90.         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, check=True)
  91.         subtitles_info = json.loads(result.stdout)
  92.         for stream in subtitles_info.get('streams', []):
  93.             if stream.get('tags', {}).get('language') == language_code:
  94.                 return True
  95.         return False
  96.     except subprocess.CalledProcessError as e:
  97.         print(f"Error al verificar subtítulos incrustados en {input_file}: {e.stderr}")
  98.         return False
  99.  
  100. def extract_embedded_subtitles(input_file, output_file):
  101.     command = [
  102.         "ffmpeg", "-i", input_file, "-map", "0:s:0", output_file
  103.     ]
  104.     try:
  105.         subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, check=True)
  106.         print(f"Subtítulos extraídos a {output_file}")
  107.     except subprocess.CalledProcessError as e:
  108.         print(f"Error al extraer subtítulos incrustados del archivo {input_file}: {e.stderr}")
  109.  
  110. def translate_srt(file_path, target_language="es"):
  111.     if not os.path.exists(file_path):
  112.         print(f"Error: No se encontraron subtítulos compatibles para el archivo {file_path}.", flush=True)
  113.         return
  114.  
  115.     with open(file_path, 'r') as file:
  116.         lines = file.readlines()
  117.  
  118.     output_file_path = file_path.replace('.en.srt', f'.{target_language}.srt').replace('.eng.srt', f'.{target_language}.srt').replace('_extracted.srt', f'.{target_language}.srt')
  119.  
  120.     with open(output_file_path, 'w') as output_file, tqdm(total=len(lines), desc="Translating", unit="subtitle") as pbar:
  121.         text_to_translate = ""
  122.         subtitle_block = []
  123.  
  124.         for line in lines:
  125.             if line.strip().isdigit():
  126.                 if subtitle_block:
  127.                     translated_text = translate_text(text_to_translate, target_language)
  128.                     output_file.writelines(subtitle_block[:-1])
  129.                     output_file.write(translated_text + '\n')
  130.                     output_file.write('\n')
  131.                     output_file.flush()
  132.                 subtitle_block = [line]
  133.                 text_to_translate = ""
  134.             elif '-->' in line:
  135.                 subtitle_block.append(line)
  136.             elif line.strip():
  137.                 text_to_translate += line.strip() + " "
  138.             else:
  139.                 subtitle_block.append(line)
  140.             pbar.update(1)
  141.  
  142.         if subtitle_block:
  143.             translated_text = translate_text(text_to_translate, target_language)
  144.             output_file.writelines(subtitle_block[:-1])
  145.             output_file.write(translated_text + '\n')
  146.             output_file.write('\n')
  147.             output_file.flush()
  148.             pbar.update(1)
  149.  
  150.     print(f'Translation completed. Output file: {output_file_path}')
  151.  
  152. def main():
  153.     parser = argparse.ArgumentParser(description="Transcodificar videos y agregar subtítulos")
  154.     parser.add_argument('input_file', type=str, help="Ruta al archivo de video a procesar")
  155.     args = parser.parse_args()
  156.  
  157.     input_file = args.input_file
  158.     temp_file = f"{os.path.splitext(input_file)[0]}_temp.mkv"
  159.  
  160.     print("Debug: Iniciando proceso principal...", flush=True)
  161.     if should_transcode(input_file):
  162.         if transcode_video(input_file, temp_file):
  163.             # Replace the original file with the transcoded one
  164.             os.replace(temp_file, input_file)
  165.  
  166.     if not has_embedded_subtitles(input_file, language_code="es"):
  167.         subtitle_file = find_subtitle_file(input_file)
  168.         if subtitle_file:
  169.             if ".es." not in subtitle_file and ".spa." not in subtitle_file:
  170.                 translate_srt(subtitle_file, target_language="es")
  171.         else:
  172.             embedded_subtitle_file = f"{os.path.splitext(input_file)[0]}.es.srt"
  173.             extract_embedded_subtitles(input_file, embedded_subtitle_file)
  174.             translate_srt(embedded_subtitle_file, target_language="es")
  175.  
  176. if __name__ == "__main__":
  177.     main()
  178.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement