Spaces:
Sleeping
Sleeping
File size: 5,237 Bytes
a5b0808 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import gradio as gr
import yaml
import re
import os
import torch
# ── STEP 1: Fix config.yaml on startup ───────────────────────
config_path = "glmocr/config.yaml"
with open(config_path, "r") as f:
config = yaml.safe_load(f)
config["pipeline"]["result_formatter"]["abandon"] = [
"number", "footnote", "aside_text", "reference",
"footer_image", "header_image",
]
config["pipeline"]["enable_layout"] = True
with open(config_path, "w") as f:
yaml.dump(config, f, default_flow_style=False, sort_keys=False)
print("✅ config.yaml fixed")
# ── STEP 2: Fix result_formatter.py on startup ───────────────
formatter_path = "glmocr/postprocess/result_formatter.py"
with open(formatter_path, "r") as f:
source = f.read()
for label in ['"header"', "'header'", '"footer"', "'footer'", '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"]:
source = re.sub(r',\s*' + re.escape(label), '', source)
source = re.sub(re.escape(label) + r'\s*,', '', source)
source = re.sub(re.escape(label), '', source)
with open(formatter_path, "w") as f:
f.write(source)
print("✅ result_formatter.py fixed")
# ── STEP 3: Load model ────────────────────────────────────────
from transformers import AutoProcessor, GlmOcrForConditionalGeneration
print("Loading model... (~2GB first run)")
processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
model = GlmOcrForConditionalGeneration.from_pretrained(
"zai-org/GLM-OCR",
torch_dtype=torch.bfloat16,
device_map="auto",
)
print("✅ Model ready on", next(model.parameters()).device)
ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
# ── STEP 4: OCR function ──────────────────────────────────────
def run_ocr(uploaded_file):
if uploaded_file is None:
return "Please upload a file.", "No regions detected."
try:
path = uploaded_file.name if hasattr(uploaded_file, 'name') else str(uploaded_file)
messages = [
{"role": "user", "content": [
{"type": "image", "url": path},
{"type": "text", "text": "Document Parsing:"}
]}
]
inputs = processor.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True,
return_dict=True, return_tensors="pt"
).to(model.device)
inputs.pop("token_type_ids", None)
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=2048)
raw = processor.decode(
output_ids[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=False
)
json_match = re.search(r'\[.*\]', raw, re.DOTALL)
regions = json.loads(json_match.group()) if json_match else []
header_count = footer_count = 0
region_lines = []
markdown_parts = []
for region in regions:
label = region.get("label", "text")
content = str(region.get("content", ""))
if label in ABANDON:
continue
if label == "header":
header_count += 1
region_lines.append("🔵 HEADER:\n" + content + "\n")
markdown_parts.append("<!-- HEADER -->\n" + content)
elif label == "footer":
footer_count += 1
region_lines.append("🟢 FOOTER:\n" + content + "\n")
markdown_parts.append("<!-- FOOTER -->\n" + content)
else:
region_lines.append("[" + label + "]: " + content[:150])
markdown_parts.append(content)
summary = (
"Headers found : " + str(header_count) + "\n"
"Footers found : " + str(footer_count) + "\n"
"Total regions : " + str(len(regions)) + "\n" + "─"*40 + "\n"
+ "\n".join(region_lines)
)
markdown = "\n\n".join(markdown_parts) if markdown_parts else raw
return markdown, summary
except Exception as e:
import traceback
return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
# ── STEP 5: Gradio UI ─────────────────────────────────────────
with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
gr.Markdown("# 🔍 GLM-OCR — Header & Footer Fix\nUpload any PDF or image. Headers 🔵 and Footers 🟢 now appear in output.")
file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
run_btn = gr.Button("▶ Run OCR", variant="primary", size="lg")
with gr.Row():
with gr.Column():
gr.Markdown("### 📄 Markdown Output")
markdown_out = gr.Textbox(lines=25, label="")
with gr.Column():
gr.Markdown("### 🗂️ Detected Regions")
regions_out = gr.Textbox(lines=25, label="")
run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
demo.launch()
|