Spaces:

leonarb
/

olmocr-demo

Running

App Files Files Community

leonarb commited on 23 days ago

Commit

b3d319d

verified ·

1 Parent(s): 0b3222e

Update app.py

Browse files

Files changed (1) hide show

app.py +18 -32

app.py CHANGED Viewed

@@ -15,23 +15,23 @@ from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
 cache_dir = "/tmp/huggingface_cache"
 os.environ["HF_HOME"] = cache_dir
 os.environ["TORCH_HOME"] = cache_dir
-os.environ["OLMOCR_LOG_PATH"] = "/tmp/olmocr-pipeline-debug.log"
 os.makedirs(cache_dir, exist_ok=True)
-# Patch logging path before olmocr import
 import logging
-original_file_handler = logging.FileHandler
-def safe_file_handler(filename, *args, **kwargs):
-    if filename == "olmocr-pipeline-debug.log":
-        filename = os.environ.get("OLMOCR_LOG_PATH", "/tmp/olmocr-pipeline-debug.log")
-    return original_file_handler(filename, *args, **kwargs)
-logging.FileHandler = safe_file_handler
-# Import olmocr pipeline after setting log path
-from olmocr.pipeline import PDFToTextOCR
-from olmocr.data.renderpdf import render_pdf_to_base64png
 from olmocr.prompts import build_finetuning_prompt
 from olmocr.prompts.anchor import get_anchor_text
 # Load model and processor
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -41,9 +41,6 @@ model = Qwen2VLForConditionalGeneration.from_pretrained(
 ).eval().to(device)
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-# Load OCR pipeline
-olmocr = PDFToTextOCR(model=model, processor=processor)
 def ocr_page(pdf_path, page_num):
     image_b64 = render_pdf_to_base64png(pdf_path, page_num + 1, target_longest_image_dim=1024)
     anchor_text = get_anchor_text(pdf_path, page_num + 1, pdf_engine="pdfreport", target_length=4000)
@@ -77,46 +74,36 @@ def ocr_page(pdf_path, page_num):
 def create_epub_from_text(text, output_path, title, author, language, cover_image):
     book = epub.EpubBook()
-    # Set metadata
     book.set_title(title)
     book.set_language(language)
     book.add_author(author)
-    # Add cover image
     with open(cover_image, "rb") as cover_file:
         cover_data = cover_file.read()
     cover_item = epub.EpubItem(uid="cover", file_name="cover.jpg", media_type="image/jpeg", content=cover_data)
     book.add_item(cover_item)
-    # Create a chapter for the content
     chapter = epub.EpubHtml(title="Content", file_name="content.xhtml", lang=language)
     chapter.set_content(f"<html><body><h1>{title}</h1><p>{text}</p></body></html>")
     book.add_item(chapter)
-    # Define Table of Contents (TOC)
     book.toc = (epub.Link("content.xhtml", "Content", "content"),)
-    # Add default NCX and OPF files
     book.add_item(epub.EpubNav())
-    # Write the EPUB file
     epub.write_epub(output_path, book)
 def convert_pdf_to_epub(pdf_file, title, author, language):
     tmp_pdf_path = pdf_file.name
-    # Read PDF to get cover
     reader = PdfReader(tmp_pdf_path)
-    first_page = reader.pages[0]
     cover_path = "/tmp/cover.jpg"
     images = convert_from_path(tmp_pdf_path, first_page=1, last_page=1)
     images[0].save(cover_path, "JPEG")
-    # Run OCR
-    ocr_text = olmocr.process(tmp_pdf_path)
-    # Write EPUB
     epub_path = "/tmp/output.epub"
     create_epub_from_text(
         text=ocr_text,
@@ -126,7 +113,6 @@ def convert_pdf_to_epub(pdf_file, title, author, language):
         language=language,
         cover_image=cover_path
     )
     return epub_path, cover_path
 def interface_fn(pdf, title, author, language):

 cache_dir = "/tmp/huggingface_cache"
 os.environ["HF_HOME"] = cache_dir
 os.environ["TORCH_HOME"] = cache_dir
 os.makedirs(cache_dir, exist_ok=True)
+# Patch logging to avoid permission errors
 import logging
+from logging import FileHandler
+class SafeFileHandler(FileHandler):
+    def __init__(self, filename, mode='a', encoding=None, delay=False, errors=None):
+        # Redirect all logs to tmp
+        safe_path = os.environ.get("OLMOCR_LOG_PATH", "/tmp/olmocr-pipeline-debug.log")
+        super().__init__(safe_path, mode, encoding, delay, errors)
+logging.FileHandler = SafeFileHandler
+# Now import olmocr
+from olmocr.run_ocr import ocr_pdf_to_text
 from olmocr.prompts import build_finetuning_prompt
 from olmocr.prompts.anchor import get_anchor_text
+from olmocr.data.renderpdf import render_pdf_to_base64png
 # Load model and processor
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ).eval().to(device)
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
 def ocr_page(pdf_path, page_num):
     image_b64 = render_pdf_to_base64png(pdf_path, page_num + 1, target_longest_image_dim=1024)
     anchor_text = get_anchor_text(pdf_path, page_num + 1, pdf_engine="pdfreport", target_length=4000)
 def create_epub_from_text(text, output_path, title, author, language, cover_image):
     book = epub.EpubBook()
     book.set_title(title)
     book.set_language(language)
     book.add_author(author)
     with open(cover_image, "rb") as cover_file:
         cover_data = cover_file.read()
     cover_item = epub.EpubItem(uid="cover", file_name="cover.jpg", media_type="image/jpeg", content=cover_data)
     book.add_item(cover_item)
     chapter = epub.EpubHtml(title="Content", file_name="content.xhtml", lang=language)
     chapter.set_content(f"<html><body><h1>{title}</h1><p>{text}</p></body></html>")
     book.add_item(chapter)
     book.toc = (epub.Link("content.xhtml", "Content", "content"),)
     book.add_item(epub.EpubNav())
     epub.write_epub(output_path, book)
 def convert_pdf_to_epub(pdf_file, title, author, language):
     tmp_pdf_path = pdf_file.name
     reader = PdfReader(tmp_pdf_path)
     cover_path = "/tmp/cover.jpg"
     images = convert_from_path(tmp_pdf_path, first_page=1, last_page=1)
     images[0].save(cover_path, "JPEG")
+    # Use official AllenAI OCR function
+    ocr_text = ocr_pdf_to_text(
+        pdf_path=tmp_pdf_path,
+        model=model,
+        processor=processor
+    )
     epub_path = "/tmp/output.epub"
     create_epub_from_text(
         text=ocr_text,
         language=language,
         cover_image=cover_path
     )
     return epub_path, cover_path
 def interface_fn(pdf, title, author, language):