Update app.py

app.py

@@ -116,16 +116,14 @@ PIPELINE_CONFIGS = {
 }
 
 @spaces.GPU
-def process_htr(image: Image.Image, document_type: Literal["letter_english", "letter_swedish", "spread_english", "spread_swedish"] = "letter_english", output_format: Literal["txt", "alto", "page", "json"] = DEFAULT_OUTPUT, custom_settings: Optional[str] = None):
+def process_htr(image_path: str, document_type: Literal["letter_english", "letter_swedish", "spread_english", "spread_swedish"] = "letter_english", output_format: Literal["txt", "alto", "page", "json"] = DEFAULT_OUTPUT, custom_settings: Optional[str] = None):
     """Process handwritten text recognition and return extracted text with specified format file."""
-    if image is None:
+    if image_path is None:
         return "Error: No image provided", None
 
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
-        image.save(temp_file.name, "PNG")
-        temp_image_path = temp_file.name
-
     try:
+        original_filename = Path(image_path).stem or "output"
+        
         if custom_settings:
             try:
                 config = json.loads(custom_settings)
@@ -134,7 +132,7 @@ def process_htr(image: Image.Image, document_type: Literal["letter_english", "le
         else:
             config = PIPELINE_CONFIGS[document_type]
 
-        collection = Collection([temp_image_path])
+        collection = Collection([image_path])
         pipeline = Pipeline.from_config(config)
         
         try:
@@ -149,7 +147,12 @@ def process_htr(image: Image.Image, document_type: Literal["letter_english", "le
         output_file_path = None
         for root, _, files in os.walk(export_dir):
             for file in files:
-                output_file_path = os.path.join(root, file)
+                old_path = os.path.join(root, file)
+                file_ext = Path(file).suffix
+                new_filename = f"{original_filename}.{output_format}" if not file_ext else f"{original_filename}{file_ext}"
+                new_path = os.path.join(root, new_filename)
+                os.rename(old_path, new_path)
+                output_file_path = new_path
                 break
         
         extracted_text = extract_text_from_collection(processed_collection)
@@ -158,12 +161,8 @@ def process_htr(image: Image.Image, document_type: Literal["letter_english", "le
         
     except Exception as e:
         return f"Error: HTR processing failed: {str(e)}", None
-    finally:
-        if os.path.exists(temp_image_path):
-            os.unlink(temp_image_path)
 
 def extract_text_from_collection(collection: Collection) -> str:
-    """Extract plain text from processed collection."""
     text_lines = []
     for page in collection.pages:
         for node in page.traverse():
@@ -175,7 +174,7 @@ def create_htrflow_mcp_server():
     demo = gr.Interface(
         fn=process_htr,
         inputs=[
-            gr.Image(type="pil", label="Upload Image"),
+            gr.Image(type="filepath", label="Upload Image or Enter URL"),
             gr.Dropdown(choices=["letter_english", "letter_swedish", "spread_english", "spread_swedish"], value="letter_english", label="Document Type"),
             gr.Dropdown(choices=CHOICES, value=DEFAULT_OUTPUT, label="Output Format"),
             gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings"),
@@ -185,7 +184,7 @@ def create_htrflow_mcp_server():
             gr.File(label="Download Output File")
         ],
         title="HTRflow MCP Server",
-        description="Process handwritten text and get extracted text with output file in specified format",
+        description="Process handwritten text from uploaded file or URL and get extracted text with output file in specified format",
         api_name="process_htr",
     )
     return demo