Update app.py

app.py

@@ -120,72 +120,49 @@ 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: 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.
-    
-    This function uses machine learning models to automatically detect, segment, and transcribe handwritten text 
-    from historical documents. It supports different document types and languages, with specialized models 
-    trained on historical handwriting from the Swedish National Archives (Riksarkivet).
-    
-    Args:
-        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 (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)
-                                - "spread_english": Two-page spread English documents with marginalia
-                                - "spread_swedish": Two-page spread Swedish documents with marginalia
-                                Default: "letter_swedish"
-        
-        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
-                                - "page": PAGE XML format with structural markup and positioning data  
-                                - "json": JSON format with structured text, layout information and metadata
-                                Default: "alto"
-        
-        custom_settings (Optional[str]): Advanced users can provide custom pipeline configuration as a 
-                                        JSON string to override the default processing steps.
-                                        Default: None (uses predefined configuration for document_type)
-        
-        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
-            - File path for direct download via gr.File (server_name/gradio_api/file=/tmp/gradio/{temp_folder}/{file_name})
-    """
+def _process_htr_pipeline(image_path: str, document_type: FormatChoices, custom_settings: Optional[str] = None) -> Collection:
+    """Process HTR pipeline and return the processed collection."""
     if not image_path:
-        error_json = json.dumps({"error": "No image provided"})
-        return error_json, None
+        raise ValueError("No image provided")
 
+    if custom_settings:
+        try:
+            config = json.loads(custom_settings)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON in custom_settings parameter")
+    else:
+        config = PIPELINE_CONFIGS[document_type]
+
+    collection = Collection([image_path])
+    pipeline = Pipeline.from_config(config)
+    
     try:
-        original_filename = Path(image_path).stem or "output"
-        
-        if custom_settings:
-            try:
-                config = json.loads(custom_settings)
-            except json.JSONDecodeError:
-                error_json = json.dumps({"error": "Invalid JSON in custom_settings parameter"})
-                return error_json, None
-        else:
-            config = PIPELINE_CONFIGS[document_type]
+        processed_collection = pipeline.run(collection)
+        return processed_collection
+    except Exception as pipeline_error:
+        raise RuntimeError(f"Pipeline execution failed: {str(pipeline_error)}")
 
-        collection = Collection([image_path])
-        pipeline = Pipeline.from_config(config)
+def htr_text(image_path: str, document_type: FormatChoices = "letter_swedish", custom_settings: Optional[str] = None) -> str:
+    """Extract text from handwritten documents using HTR."""
+    try:
+        processed_collection = _process_htr_pipeline(image_path, document_type, custom_settings)
+        extracted_text = extract_text_from_collection(processed_collection)
+        return extracted_text
         
-        try:
-            processed_collection = pipeline.run(collection)
-        except Exception as pipeline_error:
-            error_json = json.dumps({"error": f"Pipeline execution failed: {str(pipeline_error)}"})
-            return error_json, None
+    except Exception as e:
+        return f"HTR text extraction failed: {str(e)}"
 
-        extracted_text = extract_text_from_collection(processed_collection)
+def htrflow_file(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") -> str:
+    """
+    Process HTR and return a formatted file for download.
+    
+    Returns:
+        str: File path for direct download via gr.File (server_name/gradio_api/file=/tmp/gradio/{temp_folder}/{file_name})
+    """
+    try:
+        original_filename = Path(image_path).stem or "output"
+        
+        processed_collection = _process_htr_pipeline(image_path, document_type, custom_settings)
 
         temp_dir = Path(tempfile.mkdtemp())
         export_dir = temp_dir / output_format
@@ -203,23 +180,12 @@ 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()
-
-            result = {
-                "text": extracted_text,
-                "content": file_content,
-            }
-            
-            json_result = json.dumps(result, ensure_ascii=False, indent=2)
-            return json_result, output_file_path
+            return output_file_path
         else:
-            error_json = json.dumps({"error": "Failed to generate output file"})
-            return error_json, None
+            return None
         
     except Exception as e:
-        error_json = json.dumps({"error": f"HTR processing failed: {str(e)}"})
-        return error_json, None
+        return None
 
 def htrflow_visualizer(image: str, htr_document: str) -> str:
     pass
@@ -233,20 +199,34 @@ def extract_text_from_collection(collection: Collection) -> str:
     return "\n".join(text_lines)
 
 def create_htrflow_mcp_server():
-    htrflow_url = gr.Interface(
-        fn=htrflow_htr_url,
+    htr_text_interface = gr.Interface(
+        fn=htr_text,
+        inputs=[
+            gr.Image(type="filepath", label="Upload Image or Enter URL"),
+            gr.Dropdown(choices=FORMAT_CHOICES, value="letter_swedish", label="Document Type"),
+            gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings", value=""),
+        ],
+        outputs=[
+            gr.Textbox(label="Extracted Text", lines=10)
+        ],
+        description="Extract plain text from handwritten documents using HTR",
+        api_name="htr_text",
+    )
+
+    htrflow_file_interface = gr.Interface(
+        fn=htrflow_file,
         inputs=[
             gr.Image(type="filepath", label="Upload Image or Enter URL"),
             gr.Dropdown(choices=FORMAT_CHOICES, value="letter_swedish", label="Document Type"),
             gr.Dropdown(choices=FILE_CHOICES, value=DEFAULT_OUTPUT, label="Output Format"),
             gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings", value=""),
+            gr.Textbox(label="Server Name", value="https://gabriel-htrflow-mcp.hf.space", placeholder="Server URL for download links"),
         ],
         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 for file",
-        api_name="htrflow_htr_url",
+        description="Process handwritten text and get formatted file (ALTO XML, PAGE XML, JSON, or TXT)",
+        api_name="htrflow_file",
     )
 
     htrflow_viz = gr.Interface(
@@ -261,8 +241,8 @@ def create_htrflow_mcp_server():
     )
 
     demo = gr.TabbedInterface(
-        [htrflow_url, htrflow_viz],
-        ["HTR URL", "HTR Visualizer"],
+        [htr_text_interface, htrflow_file_interface, htrflow_viz],
+        ["HTR Text", "HTR File", "HTR Visualizer"],
         title="HTRflow Handwritten Text Recognition",
     )
 
@@ -270,4 +250,4 @@ def create_htrflow_mcp_server():
 
 if __name__ == "__main__":
     demo = create_htrflow_mcp_server()
-    demo.launch(mcp_server=True, share=False, debug=False)
+    demo.launch(mcp_server=True, share=False, debug=False)#