Update app.py

app.py

@@ -4,14 +4,11 @@ import tempfile
 import os
 import zipfile
 import shutil
-from typing import List, Optional, Literal, Tuple, Union, Dict
+from typing import List, Optional, Literal, Union, Dict
 from PIL import Image
 import requests
-from io import BytesIO
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-import spaces
 from pathlib import Path
+import spaces
 from visualizer import htrflow_visualizer
 from htrflow.volume.volume import Collection
 from htrflow.pipeline.pipeline import Pipeline
@@ -277,7 +274,7 @@ def _process_htr_pipeline_batch(
             results[image_name] = processed_collection
             
             if progress:
-                progress((idx + 0.9) / total_images, 
+                progress((idx + 1.0) / total_images, 
                         desc=f"Completed image {idx+1}/{total_images}: {image_name}")
             
         except Exception as e:
@@ -292,12 +289,22 @@ def _process_htr_pipeline_batch(
             pass
     
     if progress:
-        progress(1.0, desc=f"Completed processing {total_images} images!")
+        progress(1.0, desc=f"Completed processing all {total_images} images!")
     
     return results
 
 
-def htr_text_batch(
+def extract_text_from_collection(collection: Collection) -> str:
+    """Extract and combine text from all nodes in the collection."""
+    text_lines = []
+    for page in collection.pages:
+        for node in page.traverse():
+            if hasattr(node, "text") and node.text:
+                text_lines.append(node.text)
+    return "\n".join(text_lines)
+
+
+def htr_text(
     image_input: Union[str, List[str]],
     document_type: FormatChoices = "letter_swedish",
     custom_settings: Optional[str] = None,
@@ -305,22 +312,22 @@ def htr_text_batch(
     progress: gr.Progress = gr.Progress()
 ) -> str:
     """
-    Extract text from multiple handwritten documents using HTR.
-    
-    This tool processes multiple historical handwritten documents and extracts text content from each.
-    You can provide multiple image paths/URLs separated by newlines, or upload multiple files.
+    Extract text from handwritten documents using HTR.
+    Handles both single images and multiple images.
     
     Args:
         image_input: Single image path/URL, multiple paths/URLs (newline-separated), or list of uploaded files
         document_type: Type of document layout - choose based on your documents' structure and language
         custom_settings: Optional JSON configuration for advanced pipeline customization
         return_format: "separate" to show each document's text separately, "combined" to merge all text
+        progress: Progress tracker for UI updates
         
     Returns:
         Extracted text from all handwritten documents
     """
     try:
-        progress(0, desc="Starting batch HTR text extraction...")
+        if progress:
+            progress(0, desc="Starting HTR text extraction...")
         
         # Parse input to get list of images
         image_paths = parse_image_input(image_input)
@@ -328,7 +335,12 @@ def htr_text_batch(
         if not image_paths:
             return "No images provided. Please upload images or provide URLs."
         
-        progress(0.1, desc=f"Processing {len(image_paths)} images...")
+        # Adjust description based on single vs multiple
+        num_images = len(image_paths)
+        desc = f"Processing {num_images} image{'s' if num_images > 1 else ''}..."
+        
+        if progress:
+            progress(0.1, desc=desc)
         
         # Process all images
         results = _process_htr_pipeline_batch(
@@ -347,6 +359,7 @@ def htr_text_batch(
                 else:
                     all_texts.append(text)
         
+        # Return formatted result
         if return_format == "separate":
             return "\n".join(all_texts)
         else:
@@ -355,35 +368,33 @@ def htr_text_batch(
     except ValueError as e:
         return f"Input error: {str(e)}"
     except Exception as e:
-        return f"Batch HTR text extraction failed: {str(e)}"
+        return f"HTR text extraction failed: {str(e)}"
 
 
-def htrflow_file_batch(
+def htr_generate_files(
     image_input: Union[str, List[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",
     progress: gr.Progress = gr.Progress()
 ) -> str:
     """
-    Process multiple handwritten documents and generate formatted output files.
-    
-    This tool performs HTR on multiple documents and exports the results in various formats.
-    Returns a ZIP file containing all processed documents.
+    Process handwritten documents and generate formatted output files.
+    Returns a ZIP file for multiple documents, or single file for single document.
     
     Args:
         image_input: Single image path/URL, multiple paths/URLs (newline-separated), or list of uploaded files
         document_type: Type of document layout - affects segmentation and reading order
         output_format: Desired output format (txt for plain text, alto/page for XML with coordinates, json for structured data)
         custom_settings: Optional JSON configuration for advanced pipeline customization
-        server_name: Base URL of the server (used for generating download links)
+        progress: Progress tracker for UI updates
         
     Returns:
-        Path to ZIP file containing all generated files
+        Path to generated file(s)
     """
     try:
-        progress(0, desc="Starting batch HTR file processing...")
+        if progress:
+            progress(0, desc="Starting HTR file processing...")
         
         # Parse input to get list of images
         image_paths = parse_image_input(image_input)
@@ -394,14 +405,18 @@ def htrflow_file_batch(
             error_file.close()
             return error_file.name
         
-        progress(0.1, desc=f"Processing {len(image_paths)} images...")
+        num_images = len(image_paths)
+        
+        if progress:
+            progress(0.1, desc=f"Processing {num_images} image{'s' if num_images > 1 else ''}...")
         
         # Process all images
         results = _process_htr_pipeline_batch(
             image_paths, document_type, custom_settings, progress
         )
         
-        progress(0.9, desc="Creating ZIP archive...")
+        if progress:
+            progress(0.9, desc="Creating output files...")
         
         # Create temporary directory for output files
         temp_dir = Path(tempfile.mkdtemp())
@@ -434,15 +449,23 @@ def htrflow_file_batch(
                         output_files.append(new_path)
                         break
         
-        # Create ZIP file
-        zip_path = temp_dir / f"htr_batch_{output_format}.zip"
-        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            for file_path in output_files:
-                zipf.write(file_path, file_path.name)
-        
-        progress(1.0, desc=f"Batch processing complete! Processed {len(image_paths)} images.")
-        
-        return str(zip_path)
+        # Return single file or ZIP based on input count
+        if len(output_files) == 1 and len(image_paths) == 1:
+            # Single file - return directly
+            if progress:
+                progress(1.0, desc="Processing complete!")
+            return str(output_files[0])
+        else:
+            # Multiple files - create ZIP
+            zip_path = temp_dir / f"htr_output_{output_format}.zip"
+            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                for file_path in output_files:
+                    zipf.write(file_path, file_path.name)
+            
+            if progress:
+                progress(1.0, desc=f"Processing complete! Generated {len(output_files)} files.")
+            
+            return str(zip_path)
         
     except ValueError as e:
         error_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt')
@@ -451,33 +474,31 @@ def htrflow_file_batch(
         return error_file.name
     except Exception as e:
         error_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt')
-        error_file.write(f"Batch HTR file generation failed: {str(e)}")
+        error_file.write(f"HTR file generation failed: {str(e)}")
         error_file.close()
         return error_file.name
 
 
-def htrflow_visualizer_batch(
+def htr_visualize(
     image_input: Union[str, List[str]],
     htr_documents: Union[str, List[str]],
-    server_name: str = "https://gabriel-htrflow-mcp.hf.space",
     progress: gr.Progress = gr.Progress()
 ) -> str:
     """
-    Create visualizations for multiple HTR results overlaid on original documents.
-    
-    This tool generates annotated images showing detected text regions and recognized text
-    for multiple documents. Returns a ZIP file containing all visualization images.
+    Create visualizations for HTR results overlaid on original documents.
+    Returns a ZIP file for multiple documents, or single image for single document.
     
     Args:
         image_input: Original document image paths/URLs (newline-separated if string)
         htr_documents: HTR output files (ALTO/PAGE XML) - must match order of images
-        server_name: Base URL of the server (used for generating download links)
+        progress: Progress tracker for UI updates
         
     Returns:
-        Path to ZIP file containing all visualization images
+        Path to visualization file(s)
     """
     try:
-        progress(0, desc="Starting batch visualization generation...")
+        if progress:
+            progress(0, desc="Starting visualization generation...")
         
         # Parse inputs
         image_paths = parse_image_input(image_input)
@@ -492,7 +513,10 @@ def htrflow_visualizer_batch(
         if len(image_paths) != len(htr_paths):
             raise ValueError(f"Number of images ({len(image_paths)}) doesn't match number of HTR documents ({len(htr_paths)})")
         
-        progress(0.1, desc=f"Creating visualizations for {len(image_paths)} documents...")
+        num_docs = len(image_paths)
+        
+        if progress:
+            progress(0.1, desc=f"Creating visualization{'s' if num_docs > 1 else ''} for {num_docs} document{'s' if num_docs > 1 else ''}...")
         
         temp_dir = Path(tempfile.mkdtemp())
         output_files = []
@@ -502,17 +526,18 @@ def htrflow_visualizer_batch(
             try:
                 image_name = Path(image_path).stem if not image_path.startswith("http") else f"image_{idx+1}"
                 
-                progress((idx + 0.3) / len(image_paths), 
-                        desc=f"Visualizing document {idx+1}/{len(image_paths)}: {image_name}")
+                if progress:
+                    progress((idx + 0.3) / num_docs, 
+                            desc=f"Visualizing document {idx+1}/{num_docs}: {image_name}")
                 
                 # Handle image input
                 processed_image = handle_image_input(image_path, progress, 
-                                                    desc_prefix=f"[{idx+1}/{len(image_paths)}] ")
+                                                    desc_prefix=f"[{idx+1}/{num_docs}] ")
                 if processed_image.startswith(tempfile.gettempdir()):
                     temp_files.append(processed_image)
                 
                 # Generate visualization
-                viz_result = htrflow_visualizer(processed_image, htr_path, server_name)
+                viz_result = htrflow_visualizer(processed_image, htr_path, "")
                 
                 if viz_result and os.path.exists(viz_result):
                     # Move to temp dir with proper name
@@ -534,44 +559,43 @@ def htrflow_visualizer_batch(
             except:
                 pass
         
-        progress(0.9, desc="Creating ZIP archive...")
-        
-        # Create ZIP file
-        zip_path = temp_dir / "htr_visualizations.zip"
-        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            for file_path in output_files:
-                zipf.write(file_path, file_path.name)
-        
-        progress(1.0, desc=f"Visualization complete! Created {len(output_files)} visualizations.")
-        
-        return str(zip_path)
+        # Return single file or ZIP based on input count
+        if len(output_files) == 1 and num_docs == 1:
+            # Single visualization - return directly
+            if progress:
+                progress(1.0, desc="Visualization complete!")
+            return str(output_files[0])
+        else:
+            # Multiple visualizations - create ZIP
+            if progress:
+                progress(0.9, desc="Creating ZIP archive...")
+            
+            zip_path = temp_dir / "htr_visualizations.zip"
+            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                for file_path in output_files:
+                    zipf.write(file_path, file_path.name)
+            
+            if progress:
+                progress(1.0, desc=f"Visualization complete! Created {len(output_files)} visualization{'s' if len(output_files) > 1 else ''}.")
+            
+            return str(zip_path)
         
     except Exception as e:
         error_file = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt')
-        error_file.write(f"Batch visualization failed: {str(e)}")
+        error_file.write(f"Visualization failed: {str(e)}")
         error_file.close()
         return error_file.name
 
 
-def extract_text_from_collection(collection: Collection) -> str:
-    """Extract and combine text from all nodes in the collection."""
-    text_lines = []
-    for page in collection.pages:
-        for node in page.traverse():
-            if hasattr(node, "text") and node.text:
-                text_lines.append(node.text)
-    return "\n".join(text_lines)
-
-
 def create_htrflow_mcp_server():
-    # Batch HTR Text extraction interface
-    htr_text_batch_interface = gr.Interface(
-        fn=htr_text_batch,
+    # HTR Text extraction interface
+    htr_text_interface = gr.Interface(
+        fn=htr_text,
         inputs=[
             gr.Textbox(
-                label="Image Paths/URLs (one per line) or Upload Files",
-                placeholder="https://example.com/image1.jpg\nhttps://example.com/image2.jpg\n\nOR drag and drop multiple files",
-                lines=5
+                label="Image Input",
+                placeholder="Single image path/URL or multiple (one per line)\nYou can also drag and drop files here",
+                lines=3
             ),
             gr.Dropdown(
                 choices=FORMAT_CHOICES, 
@@ -593,20 +617,19 @@ def create_htrflow_mcp_server():
             ),
         ],
         outputs=[gr.Textbox(label="Extracted Text", lines=20)],
-        title="Batch Extract Text from Handwritten Documents",
-        description="Process multiple handwritten document images at once. Upload files or provide URLs (one per line)",
+        title="Extract Text from Handwritten Documents",
+        description="Process one or more handwritten document images. Works with letters and book spreads in English and Swedish.",
         api_name="htr_text_batch",
-        api_description="Extract text from multiple handwritten historical documents using advanced HTR models. Supports batch processing of letters and book spreads in English and Swedish. If a user passes a file as an input, use the upload_file_to_gradio tool, if present, to upload the file to the gradio app and create a Gradio File Input. Then use the returned path as the input to the tool",
     )
 
-    # Batch HTR File generation interface
-    htrflow_file_batch_interface = gr.Interface(
-        fn=htrflow_file_batch,
+    # HTR File generation interface
+    htr_files_interface = gr.Interface(
+        fn=htr_generate_files,
         inputs=[
             gr.Textbox(
-                label="Image Paths/URLs (one per line) or Upload Files",
-                placeholder="https://example.com/image1.jpg\nhttps://example.com/image2.jpg\n\nOR drag and drop multiple files",
-                lines=5
+                label="Image Input",
+                placeholder="Single image path/URL or multiple (one per line)\nYou can also drag and drop files here",
+                lines=3
             ),
             gr.Dropdown(
                 choices=FORMAT_CHOICES, 
@@ -626,126 +649,89 @@ def create_htrflow_mcp_server():
                 value="",
                 lines=3
             ),
-            gr.Textbox(
-                label="Server Name",
-                value="https://gabriel-htrflow-mcp.hf.space",
-                placeholder="Server URL for download links",
-                visible=False  # Hide this from UI but keep for API
-            ),
         ],
-        outputs=[gr.File(label="Download ZIP with HTR Output Files")],
-        title="Batch Generate HTR Output Files",
-        description="Process multiple handwritten documents and export in various formats. Returns a ZIP file with all results.",
+        outputs=[gr.File(label="Download HTR Output")],
+        title="Generate HTR Output Files",
+        description="Process handwritten documents and export in various formats. Returns ZIP for multiple files.",
         api_name="htrflow_file_batch",
-        api_description="Process multiple handwritten documents and generate formatted output files. Returns a ZIP containing outputs in ALTO XML (with text coordinates), PAGE XML, JSON (structured data), or plain text format. If a user passes a file as an input, use the upload_file_to_gradio tool, if present, to upload the file to the gradio app and create a Gradio File Input. Then use the returned path as the input to the tool",
     )
 
-    # Batch HTR Visualization interface
-    htrflow_viz_batch_interface = gr.Interface(
-        fn=htrflow_visualizer_batch,
+    # HTR Visualization interface
+    htr_viz_interface = gr.Interface(
+        fn=htr_visualize,
         inputs=[
             gr.Textbox(
-                label="Original Image Paths/URLs (one per line)",
-                placeholder="https://example.com/image1.jpg\nhttps://example.com/image2.jpg",
-                lines=5
+                label="Original Image Paths/URLs",
+                placeholder="One path/URL per line",
+                lines=3
             ),
             gr.File(
                 label="Upload HTR XML Files (ALTO/PAGE)", 
                 file_types=[".xml"],
                 file_count="multiple"
             ),
-            gr.Textbox(
-                label="Server Name",
-                value="https://gabriel-htrflow-mcp.hf.space",
-                placeholder="Server URL for download links",
-                visible=False  # Hide this from UI but keep for API
-            ),
         ],
-        outputs=gr.File(label="Download ZIP with Visualization Images"),
-        title="Batch Visualize HTR Results",
-        description="Create annotated images for multiple documents. Images and XML files must be in matching order.",
+        outputs=gr.File(label="Download Visualization"),
+        title="Visualize HTR Results",
+        description="Create annotated images showing detected regions and text. Files must be in matching order.",
         api_name="htrflow_visualizer_batch",
-        api_description="Generate visualization images showing HTR results overlaid on multiple original documents. Shows detected text regions, reading order, and recognized text for quality control. Returns a ZIP file with all visualizations. If a user passes a file as an input, use the upload_file_to_gradio tool, if present, to upload the file to the gradio app and create a Gradio File Input. Then use the returned path as the input to the tool",
     )
 
-    # Legacy single-file interfaces (kept for backward compatibility)
-    htr_text_interface = gr.Interface(
-        fn=lambda img, doc_type, settings: htr_text_batch(img, doc_type, settings, "separate"),
+    # Simplified interface for lambda compatibility (keeping for backward compatibility)
+    simple_text_interface = gr.Interface(
+        fn=lambda img, doc_type, settings: htr_text(img, doc_type, settings, "separate"),
         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='{"steps": [...]}',
-                value="",
-                lines=3
-            ),
+            gr.Image(type="filepath", label="Upload Image"),
+            gr.Dropdown(choices=FORMAT_CHOICES, value="letter_swedish", label="Document Type"),
+            gr.Textbox(label="Custom Settings (JSON)", value="", lines=3),
         ],
         outputs=[gr.Textbox(label="Extracted Text", lines=15)],
-        title="Extract Text (Single Document)",
-        description="Upload a single handwritten document image to extract text",
-        api_name="htr_text",
-        api_description="Extract text from handwritten historical documents using advanced HTR models. Supports letters and book spreads in English and Swedish. If a user passes a file as an input, use the upload_file_to_gradio tool, if present, to upload the file to the gradio app and create a Gradio File Input. Then use the returned path as the input to the tool",
+        api_name="_lambda_",
     )
 
-    htrflow_file_interface = gr.Interface(
-        fn=lambda img, doc_type, fmt, settings, srv: htrflow_file_batch([img], doc_type, fmt, settings, srv),
+    simple_file_interface = gr.Interface(
+        fn=lambda img, doc_type, fmt, settings, srv: htr_generate_files(img, doc_type, fmt, settings),
         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)",
-                value="",
-                lines=3
-            ),
-            gr.Textbox(
-                label="Server Name",
-                value="https://gabriel-htrflow-mcp.hf.space",
-                visible=False
-            ),
+            gr.Image(type="filepath"),
+            gr.Dropdown(choices=FORMAT_CHOICES, value="letter_swedish"),
+            gr.Dropdown(choices=FILE_CHOICES, value=DEFAULT_OUTPUT),
+            gr.Textbox(value="", lines=3),
+            gr.Textbox(value="https://gabriel-htrflow-mcp.hf.space", visible=False),
         ],
-        outputs=[gr.File(label="Download HTR Output File")],
-        title="Generate File (Single Document)",
-        description="Process a single handwritten document and export in various formats",
-        api_name="htrflow_file",
-        api_description="Process handwritten documents and generate formatted output files. Outputs can be in ALTO XML (with text coordinates), PAGE XML, JSON (structured data), or plain text format. If a user passes a file as an input, use the upload_file_to_gradio tool, if present, to upload the file to the gradio app and create a Gradio File Input. Then use the returned path as the input to the tool",
+        outputs=[gr.File()],
+        api_name="_lambda__1",
     )
 
-    # Create tabbed interface with better organization
+    # Create tabbed interface
     demo = gr.TabbedInterface(
         [
-            htr_text_batch_interface,
-            htrflow_file_batch_interface, 
-            htrflow_viz_batch_interface,
             htr_text_interface,
-            htrflow_file_interface,
+            htr_files_interface, 
+            htr_viz_interface,
+            simple_text_interface,
+            simple_file_interface,
         ],
         [
-            "📚 Batch Text Extract",
-            "📁 Batch File Generate",
-            "🖼️ Batch Visualize",
-            "📄 Single Text",
-            "💾 Single File",
+            "📚 Extract Text",
+            "📁 Generate Files",
+            "🖼️ Visualize Results",
+            "", # Hidden tabs for backward compatibility
+            "",
         ],
-        title="🖋️ HTRflow - Handwritten Text Recognition (Batch & Single Processing)",
+        title="🖋️ HTRflow - Handwritten Text Recognition",
         analytics_enabled=False,
     )
 
+    # Hide the last two tabs (for backward compatibility only)
+    demo.css = """
+        .tabitem:nth-child(4), .tabitem:nth-child(5) {
+            display: none !important;
+        }
+    """
+
     return demo
 
+
 if __name__ == "__main__":
     demo = create_htrflow_mcp_server()
     demo.launch(