112 lines
3.0 KiB
Python
112 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
OCR Text Recognition Script
|
|
Usage: python3 ocr.py <file_path> [--output <output_file>]
|
|
Supported: .pdf, .jpg, .jpeg, .png
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import json
|
|
import argparse
|
|
|
|
def extract_images_from_pdf(pdf_path):
|
|
"""Extract images from PDF"""
|
|
import fitz
|
|
|
|
images = []
|
|
doc = fitz.open(pdf_path)
|
|
|
|
for page_num in range(len(doc)):
|
|
page = doc[page_num]
|
|
page_images = page.get_images()
|
|
|
|
for img_index, img in enumerate(page_images):
|
|
xref = img[0]
|
|
base_image = doc.extract_image(xref)
|
|
image_bytes = base_image["image"]
|
|
image_ext = base_image["ext"]
|
|
|
|
output_path = f"/tmp/pdf_page{page_num+1}_img{img_index}.{image_ext}"
|
|
with open(output_path, "wb") as f:
|
|
f.write(image_bytes)
|
|
images.append(output_path)
|
|
|
|
doc.close()
|
|
return images
|
|
|
|
|
|
def ocr_file(file_path, output_path=None):
|
|
"""Perform OCR recognition on file"""
|
|
from paddleocr import PaddleOCR
|
|
|
|
# Initialize OCR
|
|
ocr = PaddleOCR(lang='ch', use_angle_cls=True)
|
|
|
|
ext = os.path.splitext(file_path)[1].lower()
|
|
|
|
if ext == '.pdf':
|
|
# PDF: extract images first, then recognize
|
|
images = extract_images_from_pdf(file_path)
|
|
all_texts = []
|
|
|
|
for img_path in images:
|
|
print(f"Recognizing: {img_path}")
|
|
result = ocr.predict(img_path)
|
|
if result:
|
|
texts = result[0].get('rec_texts', [])
|
|
all_texts.extend(texts)
|
|
|
|
# Clean up temporary images
|
|
for img_path in images:
|
|
try:
|
|
os.remove(img_path)
|
|
except:
|
|
pass
|
|
|
|
final_texts = all_texts
|
|
else:
|
|
# Image: recognize directly
|
|
print(f"Recognizing: {file_path}")
|
|
result = ocr.predict(file_path)
|
|
|
|
if result and len(result) > 0:
|
|
final_texts = result[0].get('rec_texts', [])
|
|
else:
|
|
final_texts = []
|
|
|
|
# Output results
|
|
if output_path:
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
for text in final_texts:
|
|
f.write(text + '\n')
|
|
print(f"Results saved to: {output_path}")
|
|
else:
|
|
print("\n=== OCR Recognition Results ===\n")
|
|
for text in final_texts:
|
|
print(text)
|
|
|
|
return final_texts
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='OCR Text Recognition Tool')
|
|
parser.add_argument('file', help='File to recognize (PDF or image)')
|
|
parser.add_argument('--output', '-o', help='Output file path (optional)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.exists(args.file):
|
|
print(f"Error: File not found: {args.file}")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
ocr_file(args.file, args.output)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|