Update app.py

app.py

@@ -195,6 +195,40 @@ def translate_speech(audio_file, source_language="English", target_language="Ara
         "translated_text": translation
     }
 
+# Modified document processing wrapper function
+def process_doc_wrapper(img, src, tgt):
+    if img is None:
+        # Return empty values if no image is provided
+        return None, "No document", {}, {}
+    
+    try:
+        result = process_document(img, src, tgt)
+        return (
+            result["annotated_image"],  # For doc_output (Image)
+            result["document_type"],    # For doc_type (Textbox)
+            result["extracted_text"],   # For extracted_info (JSON)
+            result["translated_text"]   # For translated_info (JSON)
+        )
+    except Exception as e:
+        print(f"Error in document processing: {e}")
+        return None, f"Error: {str(e)}", {}, {}
+
+# Modified speech translation wrapper function
+def speech_translate_wrapper(audio, src, tgt):
+    if audio is None:
+        # Return empty values if no audio is provided
+        return "No speech detected", ""
+    
+    try:
+        result = translate_speech(audio, src, tgt)
+        return (
+            result["original_text"],    # For original_text (Textbox)
+            result["translated_text"]   # For translated_text (Textbox)
+        )
+    except Exception as e:
+        print(f"Error in speech translation: {e}")
+        return f"Error: {str(e)}", ""
+
 # Gradio Interface
 def create_ui():
     with gr.Blocks(title="Police Vision Translator") as app:
@@ -218,7 +252,7 @@ def create_ui():
                     translated_info = gr.JSON(label="Translated Information")
             
             process_btn.click(
-                fn=lambda img, src, tgt: process_document(img, src, tgt),
+                fn=process_doc_wrapper,
                 inputs=[doc_input, source_lang, target_lang],
                 outputs=[doc_output, doc_type, extracted_info, translated_info]
             )
@@ -238,7 +272,7 @@ def create_ui():
                     translated_text = gr.Textbox(label="Translated Text")
             
             translate_btn.click(
-                fn=lambda audio, src, tgt: translate_speech(audio, src, tgt),
+                fn=speech_translate_wrapper,
                 inputs=[audio_input, speech_source_lang, speech_target_lang],
                 outputs=[original_text, translated_text]
             )