Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import json
 import tempfile
 import os
-from typing import List, Optional, Literal
+from typing import List, Optional, Literal, Tuple
 from PIL import Image
 import spaces
 from pathlib import Path
@@ -120,7 +120,7 @@ PIPELINE_CONFIGS = {
 }
 
 @spaces.GPU
-def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swedish", output_format: FileChoices = DEFAULT_OUTPUT, custom_settings: Optional[str] = None, server_name= "https://gabriel-htrflow-mcp.hf.space") -> str:
+def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swedish", output_format: FileChoices = DEFAULT_OUTPUT, custom_settings: Optional[str] = None, server_name: str = "https://gabriel-htrflow-mcp.hf.space") -> Tuple[str, str]:
     """
     Process handwritten text recognition (HTR) on uploaded images and return both file content and download link.
     
@@ -132,7 +132,7 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
         image_path (str): The file path or URL to the image containing handwritten text to be processed.
                          Supports common image formats like JPG, PNG, TIFF.
         
-        document_type (Literal): The type of document and language processing template to use.
+        document_type (FormatChoices): The type of document and language processing template to use.
                                 Available options:
                                 - "letter_english": Single-page English handwritten letters
                                 - "letter_swedish": Single-page Swedish handwritten letters (default)
@@ -140,7 +140,7 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
                                 - "spread_swedish": Two-page spread Swedish documents with marginalia
                                 Default: "letter_swedish"
         
-        output_format (Literal): The format for the output file containing the transcribed text.
+        output_format (FileChoices): The format for the output file containing the transcribed text.
                                 Available options:
                                 - "txt": Plain text format with line breaks
                                 - "alto": ALTO XML format with detailed layout and coordinate information
@@ -152,15 +152,24 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
                                         JSON string to override the default processing steps.
                                         Default: None (uses predefined configuration for document_type)
         
-    Returns:
-        str: JSON string containing both the file content and download link:
-             {
-                 "content": "file_content_here",
-                 "file_path": "[file_name](https://gabriel-htrflow-mcp.hf.space/gradio_api/file=/tmp/{temp-folder}/{file-name}.{file-format})"
-             }
+        server_name (str): The base URL of the server for constructing download links.
+                          Default: "https://gabriel-htrflow-mcp.hf.space"
+        
+            Returns:
+        Tuple[str, str]: A tuple containing:
+            - JSON string with extracted text, file content and download link
+            - File path for direct download via gr.File
+            
+        JSON structure:
+        {
+            "text": "extracted_plain_text_here",
+            "content": "full_file_content_here", 
+            "file_path": "[file_name](download_url)"
+        }
     """
     if not image_path:
-        return json.dumps({"error": "No image provided"})
+        error_json = json.dumps({"error": "No image provided"})
+        return error_json, None
 
     try:
         original_filename = Path(image_path).stem or "output"
@@ -169,7 +178,8 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
             try:
                 config = json.loads(custom_settings)
             except json.JSONDecodeError:
-                return json.dumps({"error": "Invalid JSON in custom_settings parameter"})
+                error_json = json.dumps({"error": "Invalid JSON in custom_settings parameter"})
+                return error_json, None
         else:
             config = PIPELINE_CONFIGS[document_type]
 
@@ -179,7 +189,10 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
         try:
             processed_collection = pipeline.run(collection)
         except Exception as pipeline_error:
-            return json.dumps({"error": f"Pipeline execution failed: {str(pipeline_error)}"})
+            error_json = json.dumps({"error": f"Pipeline execution failed: {str(pipeline_error)}"})
+            return error_json, None
+
+        extracted_text = extract_text_from_collection(processed_collection)
 
         temp_dir = Path(tempfile.mkdtemp())
         export_dir = temp_dir / output_format
@@ -197,29 +210,31 @@ def htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swed
                 break
         
         if output_file_path and os.path.exists(output_file_path):
-
             with open(output_file_path, 'r', encoding='utf-8') as f:
                 file_content = f.read()
 
-            
-            
             file_name = Path(output_file_path).name
-            download_url = f"{server_name}/gradio_api/file={output_file_path}"
+
+            server_base = server_name.rstrip('/')
+            download_url = f"{server_base}/gradio_api/file={os.path.abspath(output_file_path)}"
             markdown_link = f"[{file_name}]({download_url})"
             
             result = {
+                "text": extracted_text,
                 "content": file_content,
                 "file_path": markdown_link
             }
             
-            return json.dumps(result, ensure_ascii=False, indent=2)
+            json_result = json.dumps(result, ensure_ascii=False, indent=2)
+            return json_result, output_file_path
         else:
-            return json.dumps({"error": "Failed to generate output file"})
+            error_json = json.dumps({"error": "Failed to generate output file"})
+            return error_json, None
         
     except Exception as e:
-        return json.dumps({"error": f"HTR processing failed: {str(e)}"})
-    
-    
+        error_json = json.dumps({"error": f"HTR processing failed: {str(e)}"})
+        return error_json, None
+
 def htrflow_visualizer(image: str, htr_document: str) -> str:
     pass
 
@@ -240,7 +255,10 @@ def create_htrflow_mcp_server():
             gr.Dropdown(choices=FILE_CHOICES, value=DEFAULT_OUTPUT, label="Output Format"),
             gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings", value=""),
         ],
-        outputs=gr.Textbox(label="HTR Result (JSON)", lines=10),
+        outputs=[
+            gr.Textbox(label="HTR Result (JSON)", lines=10),
+            gr.File(label="Download HTR Output File")
+        ],
         description="Process handwritten text from uploaded file or URL and get both content and download link in JSON format",
         api_name="htrflow_htr_url",
     )