Update app.py

app.py

@@ -129,98 +129,6 @@ def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
         image = image.resize((width, height), Image.LANCZOS)
     return image
 
-def is_arabic_text(text: str) -> bool:
-    if not text:
-        return False
-    header_pattern = r'^#{1,6}\s+(.+)$'
-    paragraph_pattern = r'^(?!#{1,6}\s|!\[|```|\||\s*[-*+]\s|\s*\d+\.\s)(.+)$'
-    content_text = []
-    for line in text.split('\n'):
-        line = line.strip()
-        if not line:
-            continue
-        header_match = re.match(header_pattern, line, re.MULTILINE)
-        if header_match:
-            content_text.append(header_match.group(1))
-            continue
-        if re.match(paragraph_pattern, line, re.MULTILINE):
-            content_text.append(line)
-    if not content_text:
-        return False
-    combined_text = ' '.join(content_text)
-    arabic_chars = 0
-    total_chars = 0
-    for char in combined_text:
-        if char.isalpha():
-            total_chars += 1
-            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
-                arabic_chars += 1
-    return total_chars > 0 and (arabic_chars / total_chars) > 0.5
-
-def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
-    import base64
-    from io import BytesIO
-    markdown_lines = []
-    try:
-        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
-        for item in sorted_items:
-            category = item.get('category', '')
-            text = item.get(text_key, '')
-            bbox = item.get('bbox', [])
-            if category == 'Picture':
-                if bbox and len(bbox) == 4:
-                    try:
-                        x1, y1, x2, y2 = bbox
-                        x1, y1 = max(0, int(x1)), max(0, int(y1))
-                        x2, y2 = min(image.width, int(x2)), min(image.height, int(y2))
-                        if x2 > x1 and y2 > y1:
-                            cropped_img = image.crop((x1, y1, x2, y2))
-                            buffer = BytesIO()
-                            cropped_img.save(buffer, format='PNG')
-                            img_data = base64.b64encode(buffer.getvalue()).decode()
-                            markdown_lines.append(f"![Image](data:image/png;base64,{img_data})\n")
-                        else:
-                            markdown_lines.append("![Image](Image region detected)\n")
-                    except Exception as e:
-                        print(f"Error processing image region: {e}")
-                        markdown_lines.append("![Image](Image detected)\n")
-                else:
-                    markdown_lines.append("![Image](Image detected)\n")
-            elif not text:
-                continue
-            elif category == 'Title':
-                markdown_lines.append(f"# {text}\n")
-            elif category == 'Section-header':
-                markdown_lines.append(f"## {text}\n")
-            elif category == 'Text':
-                markdown_lines.append(f"{text}\n")
-            elif category == 'List-item':
-                markdown_lines.append(f"- {text}\n")
-            elif category == 'Table':
-                if text.strip().startswith('<'):
-                    markdown_lines.append(f"{text}\n")
-                else:
-                    markdown_lines.append(f"**Table:** {text}\n")
-            elif category == 'Formula':
-                if text.strip().startswith('$') or '\\' in text:
-                    markdown_lines.append(f"$$\n{text}\n$$\n")
-                else:
-                    markdown_lines.append(f"**Formula:** {text}\n")
-            elif category == 'Caption':
-                markdown_lines.append(f"*{text}*\n")
-            elif category == 'Footnote':
-                markdown_lines.append(f"^{text}^\n")
-            elif category in ['Page-header', 'Page-footer']:
-                continue
-            else:
-                markdown_lines.append(f"{text}\n")
-            markdown_lines.append("")
-    except Exception as e:
-        print(f"Error converting to markdown: {e}")
-        return str(layout_data)
-    return "\n".join(markdown_lines)
-
-
 @spaces.GPU
 def inference(model_name: str, image: Image.Image, text: str, max_new_tokens: int = 1024) -> str:
     try:
@@ -288,28 +196,23 @@ def process_image(
             'original_image': image,
             'raw_output': "",
             'layout_result': None,
-            'markdown_content': None
         }
         buffer = ""
         for raw_output, _ in inference(model_name, image, prompt, max_new_tokens):
             buffer = raw_output
             result['raw_output'] = buffer
             yield result
-        try:                                       
+        try:
             json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
             json_str = json_match.group(1) if json_match else buffer
-            layout_data = json.loads(json_str)                       
-                                   
+            layout_data = json.loads(json_str)
             result['layout_result'] = layout_data
-            try:
-                markdown_content = layoutjson2md(image, layout_data, text_key='text')
-                result['markdown_content'] = markdown_content
-            except Exception as e:
-                print(f"Error generating markdown: {e}")
-                result['markdown_content'] = buffer
         except json.JSONDecodeError:
-            print("Failed to parse JSON output, using raw output")
-            result['markdown_content'] = buffer
+            print("Failed to parse JSON output")
+            result['layout_result'] = {"error": "Failed to parse JSON"}
+        except Exception as e:
+            print(f"Error processing layout: {e}")
+            result['layout_result'] = {"error": str(e)}
         yield result
     except Exception as e:
         print(f"Error processing image: {e}")
@@ -317,8 +220,7 @@ def process_image(
         result = {
             'original_image': image,
             'raw_output': f"Error processing image: {str(e)}",
-            'layout_result': None,
-            'markdown_content': f"Error processing image: {str(e)}"
+            'layout_result': {"error": str(e)}
         }
         yield result
 
@@ -386,46 +288,40 @@ def create_gradio_interface():
                 with gr.Tabs():
                     with gr.Tab("📝 Extracted Content"):
                         output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=2, show_copy_button=True)
-                        with gr.Accordion("(Result.md)", open=False):
-                            markdown_output = gr.Markdown(label="Formatted Result (Result.Md)")
                     with gr.Tab("📋 Layout JSON"):
                         json_output = gr.JSON(label="Layout Analysis Results", value=None)
         def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
             try:
                 if not file_path:
-                    return "Please upload an image.", "Please upload an image.", None
+                    return "Please upload an image.", None
                 image, status = load_file_for_preview(file_path)
                 if image is None:
-                    return status, status, None
+                    return status, None
                 for result in process_image(model_name, image, min_pixels=int(min_pix) if min_pix else None, max_pixels=int(max_pix) if max_pix else None, max_new_tokens=max_tokens):
                     raw_output = result['raw_output']
-                    markdown_content = result['markdown_content'] or raw_output
-                    if is_arabic_text(markdown_content):
-                        markdown_update = gr.update(value=markdown_content, rtl=True)
-                    else:
-                        markdown_update = markdown_content
-                    yield raw_output, markdown_update, result['layout_result']
+                    layout_result = result['layout_result']
+                    yield raw_output, layout_result
             except Exception as e:
                 error_msg = f"Error processing document: {str(e)}"
                 print(error_msg)
                 traceback.print_exc()
-                yield error_msg, error_msg, None
+                yield error_msg, {"error": str(e)}
         def handle_file_upload(file_path):
             if not file_path:
                 return None, "No file loaded"
             image, page_info = load_file_for_preview(file_path)
             return image, page_info
         def clear_all():
-            return None, None, "No file loaded", "", "Click 'Process Document' to see extracted content...", None
+            return None, None, "No file loaded", None
         file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview, output])
         process_btn.click(
             process_document,
             inputs=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
-            outputs=[output, markdown_output, json_output]
+            outputs=[output, json_output]
         )
         clear_btn.click(
             clear_all,
-            outputs=[file_input, image_preview, output, markdown_output, json_output]
+            outputs=[file_input, image_preview, output, json_output]
         )
     return demo