Update app.py

app.py

@@ -15,15 +15,16 @@ from pathlib import Path
 import torch
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, Qwen2_5_VLForConditionalGeneration
 from transformers.image_utils import load_image
-
-from pdf2image import convert_from_path
+import fitz  # PyMuPDF library for PDF processing
 import html2text
 import markdown
 import tempfile
 
+# Define supported file suffixes
 pdf_suffixes = [".pdf"]
 image_suffixes = [".png", ".jpeg", ".jpg"]
 
+# LaTeX delimiter configurations
 latex_delimiters_type_a = [
     {'left': '$$', 'right': '$$', 'display': True},
     {'left': '$', 'right': '$', 'display': False},
@@ -34,6 +35,7 @@ latex_delimiters_type_b = [
 ]
 latex_delimiters_type_all = latex_delimiters_type_a + latex_delimiters_type_b
 
+# --- Model and Processor Initialization ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL_ID = "Logics-MLLM/Logics-Parsing"
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
@@ -45,6 +47,9 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 
 @spaces.GPU
 def parse_page(image: Image.Image) -> str:
+    """
+    Parses a single document page image using the Qwen2.5-VL model.
+    """
     messages = [
         {
             "role": "user",
@@ -80,23 +85,22 @@ def parse_page(image: Image.Image) -> str:
     return output_text
 
 
-def images_bytes_to_pdf_bytes(image_bytes):
-    # Memory buffer
+def images_bytes_to_pdf_bytes(image_bytes: bytes) -> bytes:
+    """
+    Converts image bytes into PDF bytes.
+    """
     pdf_buffer = BytesIO()
-
-    # Load and convert all images to RGB mode
     image = Image.open(BytesIO(image_bytes)).convert("RGB")
-
-    # Save the first image as a PDF and append the rest
     image.save(pdf_buffer, format="PDF", save_all=True)
-
-    # Get PDF bytes and reset the pointer (optional)
     pdf_bytes = pdf_buffer.getvalue()
     pdf_buffer.close()
     return pdf_bytes
 
 
-def read_fn(path):
+def read_fn(path: str or Path) -> bytes:
+    """
+    Reads a file and returns its content in bytes. Converts images to PDF bytes.
+    """
     if not isinstance(path, Path):
         path = Path(path)
     with open(str(path), "rb") as input_file:
@@ -109,110 +113,63 @@ def read_fn(path):
             raise Exception(f"Unknown file suffix: {path.suffix}")
 
 
-def safe_stem(file_path):
+def safe_stem(file_path: str) -> str:
+    """
+    Creates a safe file stem from a path.
+    """
     stem = Path(file_path).stem
-    # Keep only letters, numbers, underscores, and dots, and replace other characters with underscores
     return re.sub(r'[^\w.]', '_', stem)
 
 
-def sanitize_filename(filename: str, max_prefix_len: int = 15) -> str:
+def to_pdf(file_path: str) -> str or None:
     """
-    Sanitize filename: remove illegal characters, truncate, and add a hash to prevent duplicates.
+    Ensures the input file is in PDF format for consistent processing.
+    If the input is an image, it's converted to a temporary PDF.
     """
-    # 1. Extract the extension
-    name, ext = '', ''
-    if '.' in filename:
-        name = filename.rsplit('.', 1)[0]
-        ext = '.' + filename.rsplit('.', 1)[1].lower()
-    else:
-        name = filename
-        ext = ''
-
-    # 2. Remove illegal characters (Windows/Linux compatible)
-    # Allowed: letters, numbers, -_.()
-    name = re.sub(r'[\\/:\*\?"<>\|\s]+', '_', name)  # Replace spaces and illegal characters with underscores
-    name = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', name)  # Remove control characters
-
-    # 3. Truncate and reserve space for the hash
-    prefix = name[:max_prefix_len]
-
-    # 4. Add an MD5 prefix hash to ensure uniqueness (based on the original path or content)
-    hash_suffix = hashlib.md5(filename.encode('utf-8')).hexdigest()[:6]
-
-    # 5. Combine
-    safe_name = f"{prefix}_{hash_suffix}{ext}"
-
-    # 6. Prevent starting or ending with a dot (sensitive in some systems)
-    while safe_name.startswith('.'):
-        safe_name = safe_name[1:]
-    if len(safe_name) == 0:
-        safe_name = f"file_{hash_suffix}.bin"
-
-    if len(safe_name.encode('utf-8')) > 250:
-        # Fallback to an absolutely safe name
-        unique_hash = hashlib.md5(filename.encode('utf-8')).hexdigest()[:8]
-        safe_name = f"doc_{unique_hash}.pdf"
-
-    return safe_name
-
-
-def to_pdf(file_path):
     if file_path is None:
         return None
 
     pdf_bytes = read_fn(file_path)
-
-    # unique_filename = f'{uuid.uuid4()}.pdf'
     unique_filename = f'{safe_stem(file_path)}.pdf'
-
-    # Construct the full file path
     tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
 
-    # Write the byte data to the file
     with open(tmp_file_path, 'wb') as tmp_pdf_file:
         tmp_pdf_file.write(pdf_bytes)
 
     return tmp_file_path
 
 
-def arg_parse(ctx: 'click.Context') -> dict:
-    # Parse extra arguments
-    extra_kwargs = {}
-    i = 0
-    while i < len(ctx.args):
-        arg = ctx.args[i]
-        if arg.startswith('--'):
-            param_name = arg[2:].replace('-', '_')  # Convert parameter name format
-            i += 1
-            if i < len(ctx.args) and not ctx.args[i].startswith('--'):
-                # The parameter has a value
-                try:
-                    # Try to convert to the appropriate type
-                    if ctx.args[i].lower() == 'true':
-                        extra_kwargs[param_name] = True
-                    elif ctx.args[i].lower() == 'false':
-                        extra_kwargs[param_name] = False
-                    elif '.' in ctx.args[i]:
-                        try:
-                            extra_kwargs[param_name] = float(ctx.args[i])
-                        except ValueError:
-                            extra_kwargs[param_name] = ctx.args[i]
-                    else:
-                        try:
-                            extra_kwargs[param_name] = int(ctx.args[i])
-                        except ValueError:
-                            extra_kwargs[param_name] = ctx.args[i]
-                except:
-                    extra_kwargs[param_name] = ctx.args[i]
-            else:
-                # Boolean flag parameter
-                extra_kwargs[param_name] = True
-                i -= 1
-        i += 1
-    return extra_kwargs
-
-
-async def pdf_parse(file_path, request: gr.Request):
+def convert_pdf_to_images_fitz(pdf_path: str, dpi: int = 200) -> list:
+    """
+    Converts a PDF file to a list of PIL Images using PyMuPDF (fitz).
+    This function replaces the need for the external Poppler dependency.
+    """
+    images = []
+    logger.info(f"Converting PDF '{pdf_path}' to images with PyMuPDF.")
+    try:
+        pdf_document = fitz.open(pdf_path)
+        # Calculate zoom factor based on desired DPI
+        zoom = dpi / 72.0
+        mat = fitz.Matrix(zoom, zoom)
+
+        for page_num in range(len(pdf_document)):
+            page = pdf_document.load_page(page_num)
+            pix = page.get_pixmap(matrix=mat)
+            img_data = pix.tobytes("png")  # Use PNG for better quality
+            image = Image.open(BytesIO(img_data))
+            images.append(image)
+        pdf_document.close()
+        logger.info(f"Successfully converted {len(images)} pages.")
+    except Exception as e:
+        logger.error(f"Failed to convert PDF using PyMuPDF: {e}")
+        raise
+    return images
+
+
+async def pdf_parse(file_path: str, request: gr.Request):
+    """
+    Main parsing function that orchestrates the PDF processing pipeline.
+    """
     if file_path is None:
         logger.warning("file_path is None")
         return (
@@ -223,7 +180,9 @@ async def pdf_parse(file_path, request: gr.Request):
             None,
             "Error: No file provided"
         )
-    logger.info(f'file_path: {file_path}')
+    logger.info(f'Processing file: {file_path}')
+    
+    # Ensure file is in PDF format
     tmp_pdf_path = to_pdf(file_path)
     if tmp_pdf_path is None:
         return (
@@ -234,25 +193,36 @@ async def pdf_parse(file_path, request: gr.Request):
             None,
             "Error: Failed to process file"
         )
+
     start_time = time.time()
     try:
-        pages = convert_from_path(tmp_pdf_path, dpi=200)
+        # ** FIX: Use PyMuPDF (fitz) instead of pdf2image to avoid Poppler dependency **
+        pages = convert_pdf_to_images_fitz(tmp_pdf_path, dpi=200)
+
         html_parts = []
         for i, page in enumerate(pages):
             logger.info(f"Parsing page {i+1}/{len(pages)}")
             html = parse_page(page)
             html_parts.append(f'<div class="page-{i+1}">{html}</div>')
+        
         full_html = '\n'.join(html_parts)
         parsing_time = time.time() - start_time
+        
+        # Convert generated HTML to Markdown
         mmd = html2text.html2text(full_html)
         mmd_html = markdown.markdown(mmd)
         qwen_html = full_html
+        
+        # Create a temporary markdown file for download
         with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
             f.write(mmd)
             md_path = f.name
+            
         input_path = tmp_pdf_path
-        cost_time = f'Queue waiting time: 0, Parsing time: {parsing_time:.2f}s, Total time: {parsing_time:.2f}s'
+        cost_time = f'Parsing time: {parsing_time:.2f}s, Total time: {parsing_time:.2f}s'
+        
         return mmd_html, mmd, qwen_html, md_path, input_path, cost_time
+
     except Exception as e:
         logger.error(f"Parsing failed: {e}")
         return (
@@ -267,26 +237,10 @@ async def pdf_parse(file_path, request: gr.Request):
 
 @click.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True))
 @click.pass_context
-@click.option(
-    '--latex-delimiters-type',
-    'latex_delimiters_type',
-    type=click.Choice(['a', 'b', 'all']),
-    help="Set the type of LaTeX delimiters to use in Markdown rendering:"
-         "'a' for type '$', 'b' for type '()[]', 'all' for both types.",
-    default='all',
-)
-def main(ctx, latex_delimiters_type, **kwargs):
-    kwargs.update(arg_parse(ctx))
-    if latex_delimiters_type == 'a':
-        latex_delimiters = latex_delimiters_type_a
-    elif latex_delimiters_type == 'b':
-        latex_delimiters = latex_delimiters_type_b
-    elif latex_delimiters_type == 'all':
-        latex_delimiters = latex_delimiters_type_all
-    else:
-        raise ValueError(f"Invalid latex delimiters type: {latex_delimiters_type}.")
-
-    suffixes = pdf_suffixes + image_suffixes
+def main(ctx, **kwargs):
+    """
+    Sets up and launches the Gradio user interface.
+    """
     with gr.Blocks(head='''
         <meta name="data-spm" content="label" />
       <meta name="aplus-core" content="aplus.js" />
@@ -315,51 +269,56 @@ def main(ctx, latex_delimiters_type, **kwargs):
         })(window, document, 'script', 'aplus_queue');
       </script>
     ''') as demo:
+        gr.Markdown("# 📄 Logics-Parsing Document Analysis")
+        gr.Markdown("Upload a PDF or image file to parse its content into structured Markdown and HTML formats.")
         with gr.Row():
             with gr.Column(variant='panel', scale=5):
                 with gr.Row():
                     input_file = gr.File(label='Please upload a PDF or image (Max 20 pages for conversion)',
-                                         file_types=suffixes)
+                                         file_types=pdf_suffixes + image_suffixes)
                 with gr.Row():
-                    change_bu = gr.Button('Convert')
+                    change_bu = gr.Button('Convert', variant='primary')
                     clear_bu = gr.ClearButton(value='Clear')
                 pdf_show = PDF(label='PDF Preview', interactive=False, visible=True, height=800)
 
-                example_root = os.path.join(os.getcwd(), 'parsing/examples')
-                print(example_root)
-                logger.info(f'example_root: {example_root}')
-                if os.path.exists(example_root):
-                    with gr.Accordion('Examples:'):
-                        gr.Examples(
-                            examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
-                                      _.endswith(tuple(suffixes))],
-                            inputs=input_file
-                        )
+                example_root = 'parsing/examples'
+                logger.info(f'Looking for examples in: {example_root}')
+                if os.path.exists(example_root) and os.path.isdir(example_root):
+                    example_files = [os.path.join(example_root, f) for f in os.listdir(example_root) if f.endswith(tuple(pdf_suffixes + image_suffixes))]
+                    if example_files:
+                        with gr.Accordion('Examples:', open=True):
+                            gr.Examples(
+                                examples=example_files,
+                                inputs=input_file
+                            )
 
             with gr.Column(variant='panel', scale=5):
-                output_file = gr.File(label='Conversion Result', interactive=False)
-                cost_time = gr.Text(label='Time Cost')
+                output_file = gr.File(label='Download Markdown Result', interactive=False)
+                cost_time = gr.Text(label='Time Cost', interactive=False)
                 with gr.Tabs():
-                    with gr.Tab('MMD Rendering'):
+                    with gr.Tab('Markdown Rendering'):
                         mmd_html = gr.HTML(label='MMD Rendering')
-                    # with gr.Tab('mmd html text'):
-                    #     mmd_html_text = gr.TextArea(lines=45, show_copy_button=True)
-                    with gr.Tab('MMD'):
-                        mmd = gr.TextArea(lines=45, show_copy_button=True)
-                    with gr.Tab('Qwen HTML'):
-                        raw_html = gr.TextArea(lines=45, show_copy_button=True)
-
-        clear_bu.add([input_file, pdf_show, mmd, raw_html, output_file, mmd_html, cost_time])
-        cna = gr.Textbox(visible=False)
-        input_file.change(fn=to_pdf, inputs=input_file, outputs=pdf_show)
+                    with gr.Tab('Markdown Source'):
+                        mmd = gr.TextArea(lines=45, show_copy_button=True, label="Markdown Source")
+                    with gr.Tab('Generated HTML'):
+                        raw_html = gr.TextArea(lines=45, show_copy_button=True, label="Generated HTML")
+
+        # Define component list for clearing
+        components_to_clear = [input_file, pdf_show, mmd, raw_html, output_file, mmd_html, cost_time]
+        
+        clear_bu.add(components_to_clear)
+        
+        input_file.change(fn=to_pdf, inputs=input_file, outputs=pdf_show, show_progress="full")
+        
         change_bu.click(
             fn=pdf_parse,
             inputs=[input_file],
             outputs=[mmd_html, mmd, raw_html, output_file, pdf_show, cost_time],
-            concurrency_limit=15
+            concurrency_limit=15,
+            show_progress="full"
         )
 
-    demo.launch()
+    demo.launch(debug=True)
 
 
 if __name__ == '__main__':