Update app.py

app.py

@@ -129,6 +129,97 @@ 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:
@@ -182,47 +273,33 @@ def inference(model_name: str, image: Image.Image, text: str, max_new_tokens: in
         traceback.print_exc()
         yield f"Error during inference: {str(e)}", f"Error during inference: {str(e)}"
 
+@spaces.GPU
 def process_image(
     model_name: str,
     image: Image.Image,
     min_pixels: Optional[int] = None,
     max_pixels: Optional[int] = None,
     max_new_tokens: int = 1024
-) -> Dict[str, Any]:
+):
     try:
         if min_pixels or max_pixels:
             image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
-        result = {
-            'original_image': image,
-            'raw_output': "",
-            'layout_result': None,
-        }
         buffer = ""
         for raw_output, _ in inference(model_name, image, prompt, max_new_tokens):
             buffer = raw_output
-            result['raw_output'] = buffer
-            yield result
+            yield buffer, None  # Yield raw OCR stream and None for JSON during processing
         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)
-            result['layout_result'] = layout_data
+            yield buffer, layout_data  # Final yield with raw OCR and parsed JSON
         except json.JSONDecodeError:
-            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
+            print("Failed to parse JSON output, using raw output")
+            yield buffer, None  # If JSON parsing fails, yield raw OCR with no JSON
     except Exception as e:
         print(f"Error processing image: {e}")
         traceback.print_exc()
-        result = {
-            'original_image': image,
-            'raw_output': f"Error processing image: {str(e)}",
-            'layout_result': {"error": str(e)}
-        }
-        yield result
+        yield f"Error processing image: {str(e)}", None
 
 def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
     if not file_path or not os.path.exists(file_path):
@@ -287,9 +364,9 @@ def create_gradio_interface():
             with gr.Column(scale=2):
                 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.Tab("📋 Layout JSON"):
-                        json_output = gr.JSON(label="Layout Analysis Results", value=None)
+                        raw_output = gr.Textbox(label="Raw OCR Output", interactive=False, lines=10, show_copy_button=True)
+                    with gr.Tab("📋 Layout Analysis Results"):
+                        json_output = gr.JSON(label="Layout Analysis Results")
         def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
             try:
                 if not file_path:
@@ -297,31 +374,29 @@ def create_gradio_interface():
                 image, status = load_file_for_preview(file_path)
                 if image is 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']
-                    layout_result = result['layout_result']
-                    yield raw_output, layout_result
+                for raw_buffer, layout_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):
+                    yield raw_buffer, layout_result
             except Exception as e:
                 error_msg = f"Error processing document: {str(e)}"
                 print(error_msg)
                 traceback.print_exc()
-                yield error_msg, {"error": str(e)}
+                yield error_msg, None
         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
+                return None
+            image, status = load_file_for_preview(file_path)
+            return image
         def clear_all():
-            return None, None, "No file loaded", None
-        file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview, output])
+            return None, None, "", None
+        file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview])
         process_btn.click(
             process_document,
             inputs=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
-            outputs=[output, json_output]
+            outputs=[raw_output, json_output]
         )
         clear_btn.click(
             clear_all,
-            outputs=[file_input, image_preview, output, json_output]
+            outputs=[file_input, image_preview, raw_output, json_output]
         )
     return demo