Update app.py

app.py

@@ -125,22 +125,47 @@ def transcribe_video(video_path):
         }
         for segment in result["segments"]
     ]
-    return transcript_with_timestamps
+    # Get the detected language
+    detected_language = result["language"]
+    logger.debug(f"Detected language:\n{detected_language}")
+    return transcript_with_timestamps, detected_language
 
 # Function to get the appropriate translation model based on target language
-def get_translation_model(target_language):
-    # Map of target languages to their corresponding model names
-    model_map = {
-        "es": "Helsinki-NLP/opus-mt-en-es",  # English to Spanish
-        "fr": "Helsinki-NLP/opus-mt-en-fr",  # English to French
-        "zh": "Helsinki-NLP/opus-mt-en-zh",  # English to Chinese
-        # Add more languages as needed
-    }
-    return model_map.get(target_language, "Helsinki-NLP/opus-mt-en-zh")  # Default to Chinese if not found
+def get_translation_model(source_language, target_language):
+    """
+    Get the translation model based on the source and target language.
+
+    Parameters:
+    - target_language (str): The language to translate the content into (e.g., 'es', 'fr').
+    - source_language (str): The language of the input content (default is 'en' for English).
+    
+    Returns:
+    - str: The translation model identifier.
+    """
+    # List of allowable languages
+    allowable_languages = ["en", "es", "fr", "zh", "de", "it", "pt", "ja", "ko", "ru"]
+
+    # Validate source and target languages
+    if source_language not in allowable_languages:
+        logger.debug(f"Invalid source language '{source_language}'. Supported languages are: {', '.join(allowable_languages)}")
+        # Return a default model if source language is invalid
+        source_language = "en"  # Default to 'en'
+
+    if target_language not in allowable_languages:
+        logger.debug(f"Invalid target language '{target_language}'. Supported languages are: {', '.join(allowable_languages)}")
+        # Return a default model if target language is invalid
+        target_language = "zh"  # Default to 'zh'
+
+    if source_language == target_language:
+        source_language = "en"  # Default to 'en'
+        target_language = "zh"  # Default to 'zh'
+
+    # Return the model using string concatenation
+    return f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
 
-def translate_text(transcription_json, target_language):
+def translate_text(transcription_json, source_language, target_language):
     # Load the translation model for the specified target language
-    translation_model_id = get_translation_model(target_language)
+    translation_model_id = get_translation_model(source_language, target_language)
     logger.debug(f"Translation model: {translation_model_id}")
     translator = pipeline("translation", model=translation_model_id)
 
@@ -271,7 +296,7 @@ def replace_audio_in_video(video_path: str, new_audio_path: str, final_video_pat
     except Exception as e:
         logger.error(f"Error replacing audio in video: {e}")
 
-def upload_and_manage(file, language, mode="transcription"):
+def upload_and_manage(file, target_language, mode="transcription"):
     if file is None:
         return None, [], None, "No file uploaded. Please upload a video/audio file."
 
@@ -286,17 +311,17 @@ def upload_and_manage(file, language, mode="transcription"):
         list_available_fonts()
 
         # Step 1: Transcribe audio from uploaded media file and get timestamps
-        transcription_json = transcribe_video(file.name)
+        transcription_json, source_language = transcribe_video(file.name)
 
         # Step 2: Translate the transcription
-        translated_json = translate_text(transcription_json, language)
+        translated_json = translate_text(transcription_json, source_language, target_language)
 
         # Step 3: Add transcript to video based on timestamps
         add_transcript_to_video(file.name, translated_json, output_video_path)
 
         # Step 4 (Optional): Generate voiceover if mode is "transcription_voiceover"
         if mode == "Transcription with Voiceover":
-            generate_voiceover(translated_json, language, voiceover_path)
+            generate_voiceover(translated_json, target_language, voiceover_path)
             replace_audio_in_video(output_video_path, voiceover_path, output_video_path)
 
         # Convert translated JSON into a format for the editable table