mirror of
https://github.com/lemeow125/DocManagerBackend.git
synced 2025-01-18 17:13:00 +08:00
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
from io import BytesIO
|
|
from documents.models import Document
|
|
from django.db.models.signals import post_save
|
|
from django.dispatch import receiver
|
|
from config.settings import MEDIA_ROOT
|
|
import os
|
|
import fitz
|
|
import pytesseract
|
|
from PIL import Image
|
|
from .models import Document
|
|
|
|
|
|
@receiver(post_save, sender=Document)
|
|
def document_post_save(sender, instance, **kwargs):
|
|
if not instance.ocr_metadata:
|
|
metadata = ""
|
|
with fitz.open(os.path.join(MEDIA_ROOT, instance.file.name)) as doc:
|
|
mat = fitz.Matrix(1.2, 1.2)
|
|
for page in doc:
|
|
pix = page.get_pixmap(matrix=mat)
|
|
# Convert pixmap to bytes
|
|
img_bytes = pix.tobytes()
|
|
|
|
# Create a BytesIO object
|
|
img_buffer = BytesIO(img_bytes)
|
|
|
|
# Create a PIL Image object from the bytes
|
|
img = Image.open(img_buffer)
|
|
|
|
# Perform OCR
|
|
text = pytesseract.image_to_string(img).strip()
|
|
metadata += text
|
|
|
|
instance.ocr_metadata = metadata
|
|
instance.save()
|