')
+ print(f"检测到 {mol_count} 个分子结构")
+ print(f"检测到 {mol_table_count} 个分子表格")
+
+ return result
+
+ except Exception as e:
+ print(f"处理PDF时出错: {e}")
+ raise
+
+
+def main():
+ """主函数 - 示例用法"""
+
+ # 模型文件路径配置（请根据实际情况修改）
+ model_paths = {
+ # "mol_model_path": "/path/to/your/mol_detection_model",
+ # "table_model_path": "/path/to/your/table_detection_model",
+ }
+
+ # 示例配置
+ processor_config = {
+ "device": "cuda", # 或 "cpu"
+ "debug": True, # 启用调试模式
+ "with_mol_detect": True,
+ "with_table_detect": True,
+ # 可以添加更多配置项
+ "use_yolo_mol_model": True,
+ "use_yolo_table_model": True,
+ "use_yolo_table_model_v2": True,
+ }
+
+ # 示例PDF路径（请替换为实际路径）
+ pdf_path = "example_chemistry_paper.pdf"
+
+ if os.path.exists(pdf_path):
+ try:
+ result = process_pdf_with_molecule_detection(
+ pdf_path=pdf_path,
+ processor_config=processor_config,
+ model_paths=model_paths
+ )
+
+ print("\n分子识别成功完成！")
+ print("输出中包含以下标签：")
+ print("- ...: 化学分子结构")
+ print("- ...: 包含分子数据的表格")
+ print("\n特点：")
+ print("- IOU>90%的分子表格会替换原有的普通表格")
+ print("- 保持原有的文档结构和顺序")
+ print("- 支持各种分子结构和表格格式")
+
+ except Exception as e:
+ print(f"错误: {e}")
+ else:
+ print(f"示例PDF文件不存在: {pdf_path}")
+ print("请将此脚本中的pdf_path变量设置为实际的PDF文件路径")
+ print("\n使用方法:")
+ print("1. 确保安装了img2mol及其依赖")
+ print("2. 配置模型文件路径（如果需要）")
+ print("3. 设置正确的PDF文件路径")
+ print("4. 运行脚本")
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/marker/examples/real_molecule_example.py b/marker/examples/real_molecule_example.py
new file mode 100644
index 00000000..5f62fbd0
--- /dev/null
+++ b/marker/examples/real_molecule_example.py
@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+"""
+真实img2mol集成使用示例
+展示如何正确配置和使用真实的分子识别功能
+"""
+
+import os
+import sys
+from pathlib import Path
+
+# 添加marker路径
+current_dir = Path(__file__).parent.parent.parent
+sys.path.append(str(current_dir))
+
+from marker.main import marker_main
+from marker.config.parser import ConfigParser
+from marker.config.processor import apply_config
+
+def create_real_molecule_config():
+ """创建真实img2mol处理的配置"""
+
+ # img2mol处理器配置
+ processor_config = {
+ # 基础设置
+ 'device': 'cuda' if torch.cuda.is_available() else 'cpu', # 优先使用GPU
+ 'debug': True, # 开启调试模式
+ 'num_workers': 1, # 工作线程数
+ 'padding': 0, # 图像填充
+
+ # 分子检测设置
+ 'with_mol_detect': True, # 启用分子检测
+ 'use_yolo_mol_model': True, # 使用YOLO分子模型
+ 'new_class_token': True, # 使用新的类标记
+
+ # 表格检测设置
+ 'with_table_detect': True, # 启用表格检测
+ 'use_yolo_table_model': True, # 使用YOLO表格模型
+ 'use_yolo_table_model_v2': True, # 使用YOLO表格模型v2
+
+ # OCR设置
+ 'use_trocr_mfr_model_v3': True, # 使用TrOCR数学公式识别模型v3
+ 'use_got_ocr_model': True, # 使用GOT OCR模型
+ 'with_ocr': False, # 在分子检测时不使用OCR
+
+ # 模型预加载设置（可选，影响启动速度）
+ 'preload_table_and_ocr_model': True, # 预加载表格和OCR模型
+
+ # 高级设置
+ 'use_tta': True, # 使用测试时间增强
+ 'coref': True, # 使用共指消解
+ 'with_padding': False, # 不使用图像填充
+
+ # 模型路径（如果需要指定特定模型）
+ # 'MolDetect_mol_path': '/path/to/mol/model',
+ # 'model_dir': '/path/to/models'
+ }
+
+ # marker主要配置
+ config = {
+ # 启用分子检测
+ 'use_molecule_detection': True,
+
+ # 传递处理器配置
+ 'processor_config': processor_config,
+
+ # 其他marker配置
+ 'max_pages': None, # 处理所有页面
+ 'start_page': None, # 从第一页开始
+ 'languages': ['en'], # 支持的语言
+ 'batch_multiplier': 1, # 批处理倍数
+ 'ocr_all_pages': False, # 不对所有页面进行OCR
+
+ # 输出格式
+ 'output_format': 'markdown', # 输出格式为markdown
+ 'extract_images': True, # 提取图像
+ }
+
+ return config
+
+def process_chemical_pdf(pdf_path, output_dir=None):
+ """
+ 处理化学PDF文档，检测分子和表格
+
+ Args:
+ pdf_path: PDF文件路径
+ output_dir: 输出目录，如果为None则使用默认目录
+
+ Returns:
+ 处理结果的路径
+ """
+
+ # 验证输入文件
+ pdf_path = Path(pdf_path)
+ if not pdf_path.exists():
+ raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")
+
+ # 设置输出目录
+ if output_dir is None:
+ output_dir = pdf_path.parent / f"{pdf_path.stem}_molecule_output"
+ else:
+ output_dir = Path(output_dir)
+
+ output_dir.mkdir(exist_ok=True)
+
+ print(f"🧬 开始处理化学PDF: {pdf_path}")
+ print(f"📁 输出目录: {output_dir}")
+
+ # 创建配置
+ config = create_real_molecule_config()
+
+ try:
+ # 调用marker主函数
+ result = marker_main(
+ pdf_path=str(pdf_path),
+ output_dir=str(output_dir),
+ config=config,
+ artifact_dict={
+ 'processor_config': config['processor_config'] # 确保传递处理器配置
+ }
+ )
+
+ print(f"✅ 处理完成!")
+ print(f"📄 输出文件: {result}")
+
+ return result
+
+ except Exception as e:
+ print(f"❌ 处理失败: {e}")
+ import traceback
+ traceback.print_exc()
+ return None
+
+def main():
+ """主函数"""
+ import argparse
+
+ parser = argparse.ArgumentParser(description='使用真实img2mol处理化学PDF文档')
+ parser.add_argument('pdf_path', help='输入PDF文件路径')
+ parser.add_argument('--output-dir', '-o', help='输出目录路径')
+ parser.add_argument('--mock', action='store_true', help='使用mock模式（测试用）')
+
+ args = parser.parse_args()
+
+ # 如果使用mock模式，修改配置
+ if args.mock:
+ print("🎭 使用Mock模式进行测试")
+ config = create_real_molecule_config()
+ config['processor_config']['use_mock_data'] = True
+
+ # 处理PDF
+ result = process_chemical_pdf(args.pdf_path, args.output_dir)
+
+ if result:
+ print(f"\n🎉 处理成功完成!")
+ print(f"📋 结果文件: {result}")
+ print(f"\n💡 提示: 查看输出的markdown文件，其中包含:")
+ print(f" • ... 标签标记的分子结构")
+ print(f" • ... 标签标记的分子表格")
+ else:
+ print(f"\n❌ 处理失败，请检查错误消息")
+ sys.exit(1)
+
+if __name__ == "__main__":
+ # 导入必要的模块
+ try:
+ import torch
+ except ImportError:
+ print("Warning: PyTorch not available, defaulting to CPU")
+ torch = None
+
+ main()
\ No newline at end of file
diff --git a/marker/oss_uploader.py b/marker/oss_uploader.py
new file mode 100644
index 00000000..6d238a68
--- /dev/null
+++ b/marker/oss_uploader.py
@@ -0,0 +1,82 @@
+
+import os
+import boto3
+from botocore.config import Config
+import string
+import random
+from datetime import datetime
+from pathlib import Path
+import traceback
+from typing import Union
+from dotenv import load_dotenv
+
+load_dotenv(os.path.join(os.path.dirname(__file__), '..', '.env'))
+
+
+class S3Client:
+ def __init__(self,
+ endpoint_url=os.getenv("AWS_S3_ENDPOINT_URL"),
+ aws_access_key_id=os.getenv("AWS_S3_ACCESS_KEY_ID"),
+ aws_secret_access_key=os.getenv("AWS_S3_SECRET_ACCESS_KEY"),
+ region_name=os.getenv("AWS_S3_REGION_NAME"),
+ bucket=os.getenv("AWS_STORAGE_BUCKET_NAME")
+ ):
+ self.connection_data = {
+ 'endpoint_url': endpoint_url,
+ 'aws_access_key_id': aws_access_key_id,
+ 'aws_secret_access_key': aws_secret_access_key,
+ 'region_name': region_name,
+ }
+ self.bucket = bucket
+ self.s3_session = boto3.client('s3', config=Config(s3={'addressing_style': 'virtual'}), **self.connection_data)
+
+ @staticmethod
+ def random_str():
+ alphabet = string.ascii_lowercase + string.digits
+ return ''.join(random.choices(alphabet, k=6))
+
+ def s3_upload_from_file(self, filename: str, filebytes: bytes):
+ p = Path(filename)
+ date = datetime.now().strftime('%Y%m%d')
+ key = f'qilu-brain/wemol/{date}/{p.stem}_qilu-brain_{self.random_str()}{p.suffix}'
+ return self._s3_upload(filebytes, key)
+
+ def download_from_s3(self, key: str):
+ try:
+ response = self.s3_session.get_object(Bucket=self.bucket, Key=key)
+ # Extract and read the file content from the streaming body
+ content = response['Body'].read()
+ return content
+ except Exception as e:
+ traceback.print_exc()
+ print(f'download from s3 failed, cause {str(e)}')
+ return None
+
+ def delete_from_s3(self, key: str):
+ try:
+ self.s3_session.delete_object(Bucket=self.bucket, Key=key)
+ return True
+ except Exception as e:
+ traceback.print_exc()
+ print(f'delete from s3 failed, cause {str(e)}')
+ return False
+
+ def _s3_upload(self, body: Union[str, bytes], key: str):
+ try:
+ _ = self.s3_session.put_object(Body=body, Bucket=self.bucket, Key=key)
+ # 生成url，不过期的
+ url = self.s3_session.generate_presigned_url(
+ ClientMethod='get_object',
+ Params={'Bucket': self.bucket, 'Key': key},
+ ExpiresIn=None,
+ HttpMethod='GET'
+ )
+ return {
+ 'url': url,
+ 'key': key
+ }
+
+ except Exception as e:
+ traceback.print_exc()
+ print(f'upload object failed, cause {str(e)}')
+ return {}
diff --git a/marker/providers/__init__.py b/marker/providers/__init__.py
index b8c48373..a89ced82 100644
--- a/marker/providers/__init__.py
+++ b/marker/providers/__init__.py
@@ -75,19 +75,40 @@ def __enter__(self):
def get_font_css():
from weasyprint import CSS
from weasyprint.text.fonts import FontConfiguration
+ import os
font_config = FontConfiguration()
- css = CSS(string=f'''
- @font-face {{
- font-family: GoNotoCurrent-Regular;
- src: url({settings.FONT_PATH});
- font-display: swap;
- }}
+
+ # Check if font file exists and create proper file URI
+ font_src = ""
+ if os.path.exists(settings.FONT_PATH):
+ # Convert to absolute path and use file:// URI scheme
+ abs_font_path = os.path.abspath(settings.FONT_PATH)
+ font_src = f"url(file://{abs_font_path})"
+
+ # Create CSS with proper font fallback
+ css_content = f'''
body {{
- font-family: {settings.FONT_NAME.split(".")[0]}, sans-serif;
+ font-family: "DejaVu Sans", "Liberation Sans", Arial, sans-serif;
font-variant-ligatures: none;
font-feature-settings: "liga" 0;
- text-rendering: optimizeLegibility;
}}
- ''', font_config=font_config)
+ '''
+
+ # Only add @font-face if font file exists
+ if font_src:
+ css_content = f'''
+ @font-face {{
+ font-family: GoNotoCurrent-Regular;
+ src: {font_src};
+ font-display: swap;
+ }}
+ body {{
+ font-family: GoNotoCurrent-Regular, "DejaVu Sans", "Liberation Sans", Arial, sans-serif;
+ font-variant-ligatures: none;
+ font-feature-settings: "liga" 0;
+ }}
+ '''
+
+ css = CSS(string=css_content, font_config=font_config)
return css
diff --git a/marker/providers/document.py b/marker/providers/document.py
index 4dfd1038..61a33c58 100644
--- a/marker/providers/document.py
+++ b/marker/providers/document.py
@@ -23,7 +23,7 @@
div, p {
max-width: 100%;
- word-break: break-word;
+ overflow-wrap: break-word;
font-size: 10pt;
}
@@ -62,23 +62,165 @@ def __init__(self, filepath: str, config=None):
super().__init__(self.temp_pdf_path, config)
def __del__(self):
- if os.path.exists(self.temp_pdf_path):
- os.remove(self.temp_pdf_path)
+ # if os.path.exists(self.temp_pdf_path):
+ # os.remove(self.temp_pdf_path)
+ pass
+
+ def _convert_image_to_base64(self, image):
+ """Convert mammoth image to base64 data URI"""
+ try:
+ with image.open() as image_bytes:
+ import base64
+ return "data:" + image.content_type + ";base64," + base64.b64encode(image_bytes.read()).decode()
+ except Exception as e:
+ print(f"Failed to convert image: {e}")
+ return ""
def convert_docx_to_pdf(self, filepath: str):
from weasyprint import CSS, HTML
import mammoth
+ import re
with open(filepath, "rb") as docx_file:
- # we convert the docx to HTML
- result = mammoth.convert_to_html(docx_file)
+ # Configure style mapping to preserve heading levels
+ style_map = """
+ p[style-name='Heading 1'] => h1:fresh
+ p[style-name='Heading 2'] => h2:fresh
+ p[style-name='Heading 3'] => h3:fresh
+ p[style-name='Heading 4'] => h4:fresh
+ p[style-name='Heading 5'] => h5:fresh
+ p[style-name='Heading 6'] => h6:fresh
+ p[style-name='标题 1'] => h1:fresh
+ p[style-name='标题 2'] => h2:fresh
+ p[style-name='标题 3'] => h3:fresh
+ p[style-name='标题 4'] => h4:fresh
+ p[style-name='标题 5'] => h5:fresh
+ p[style-name='标题 6'] => h6:fresh
+ p:empty => p:fresh
+ """
+
+ # Configure mammoth options for better conversion
+ convert_options = {
+ "convert_image": mammoth.images.img_element(lambda image: {
+ "src": self._convert_image_to_base64(image)
+ }),
+ "ignore_empty_paragraphs": False,
+ "style_map": style_map
+ }
+
+ # we convert the docx to HTML with better options
+ result = mammoth.convert_to_html(docx_file, **convert_options)
html = result.value
+
+ # Print conversion messages if any
+ if result.messages:
+ print(f"Mammoth conversion messages: {result.messages}", flush=True)
+
+ # Post-process HTML to normalize headings and preserve empty lines
+ html = self._normalize_html(html)
+
+ # Debug: Print HTML content length and preview
+ print(f"Generated HTML length: {len(html)} characters", flush=True)
+ if html:
+ preview = html[:500].replace('\n', ' ')
+ print(f"HTML preview: {preview}...", flush=True)
+ else:
+ print("WARNING: Generated HTML is empty!", flush=True)
# We convert the HTML into a PDF
- HTML(string=self._preprocess_base64_images(html)).write_pdf(
- self.temp_pdf_path,
- stylesheets=[CSS(string=css), self.get_font_css()]
- )
+ processed_html = self._preprocess_base64_images(html)
+ if not processed_html.strip():
+ print("ERROR: Processed HTML is empty, adding fallback content", flush=True)
+ processed_html = "Document conversion failed - no content extracted
"
+
+ # Wrap in proper HTML structure for better PDF conversion
+ if not processed_html.startswith('{processed_html}