Update app.py

app.py

@@ -481,7 +481,7 @@ def htr_generate_files(
 
 def htr_visualize(
     image_input: Union[str, List[str]],
-    htr_documents: Union[str, List[str]],
+    htr_documents: Union[List[str], None],
     progress: gr.Progress = gr.Progress()
 ) -> str:
     """
@@ -502,14 +502,33 @@ def htr_visualize(
         
         # Parse inputs
         image_paths = parse_image_input(image_input)
-        htr_paths = parse_image_input(htr_documents) if isinstance(htr_documents, str) else htr_documents
+        
+        # Handle htr_documents - it should be a list of file paths
+        if not htr_documents:
+            raise ValueError("No HTR documents provided")
+        
+        # If htr_documents is a list of file objects from Gradio, extract paths
+        htr_paths = []
+        if isinstance(htr_documents, list):
+            for doc in htr_documents:
+                if isinstance(doc, str):
+                    htr_paths.append(doc)
+                elif hasattr(doc, 'name'):
+                    htr_paths.append(doc.name)
+                else:
+                    htr_paths.append(str(doc))
+        else:
+            # Single file case
+            if isinstance(htr_documents, str):
+                htr_paths = [htr_documents]
+            elif hasattr(htr_documents, 'name'):
+                htr_paths = [htr_documents.name]
+            else:
+                htr_paths = [str(htr_documents)]
         
         if not image_paths:
             raise ValueError("No images provided")
         
-        if not htr_paths:
-            raise ValueError("No HTR documents provided")
-        
         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)})")
         
@@ -536,14 +555,20 @@ def htr_visualize(
                 if processed_image.startswith(tempfile.gettempdir()):
                     temp_files.append(processed_image)
                 
-                # Generate visualization
-                viz_result = htrflow_visualizer(processed_image, htr_path, "")
+                # Generate visualization - use the last parameter for output path
+                output_viz_path = str(temp_dir / f"{image_name}_visualization.png")
+                viz_result = htrflow_visualizer(processed_image, htr_path, output_viz_path)
                 
-                if viz_result and os.path.exists(viz_result):
-                    # Move to temp dir with proper name
+                # Check if visualization was created
+                if os.path.exists(output_viz_path):
+                    output_files.append(Path(output_viz_path))
+                elif viz_result and os.path.exists(viz_result):
+                    # Fallback: if viz_result points to a different file
                     viz_path = temp_dir / f"{image_name}_visualization.png"
                     shutil.move(viz_result, str(viz_path))
                     output_files.append(viz_path)
+                else:
+                    raise ValueError("Visualization generation failed - no output file created")
                 
             except Exception as e:
                 # Create error file for this visualization
@@ -551,6 +576,7 @@ def htr_visualize(
                 with open(error_path, 'w') as f:
                     f.write(f"Visualization failed: {str(e)}")
                 output_files.append(error_path)
+                print(f"Error visualizing {image_name}: {str(e)}")
         
         # Cleanup temp files
         for temp_file in temp_files:
@@ -619,7 +645,7 @@ def create_htrflow_mcp_server():
         outputs=[gr.Textbox(label="Extracted Text", lines=20)],
         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_name="htr_text",
     )
 
     # HTR File generation interface
@@ -653,7 +679,7 @@ def create_htrflow_mcp_server():
         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_name="htr_generate_files",
     )
 
     # HTR Visualization interface
@@ -674,32 +700,7 @@ def create_htrflow_mcp_server():
         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",
-    )
-
-    # 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"),
-            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)],
-        api_name="_lambda_",
-    )
-
-    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"),
-            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()],
-        api_name="_lambda__1",
+        api_name="htr_visualize",
     )
 
     # Create tabbed interface
@@ -708,27 +709,16 @@ def create_htrflow_mcp_server():
             htr_text_interface,
             htr_files_interface, 
             htr_viz_interface,
-            simple_text_interface,
-            simple_file_interface,
         ],
         [
             "📚 Extract Text",
             "📁 Generate Files",
             "🖼️ Visualize Results",
-            "", # Hidden tabs for backward compatibility
-            "",
         ],
         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