mirror of
https://github.com/lemeow125/DocManagerBackend.git
synced 2025-01-19 09:33:01 +08:00
117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
from django.core.management.base import BaseCommand, CommandError
|
|
|
|
from io import BytesIO
|
|
from documents.models import Document
|
|
from PIL import Image
|
|
import pytesseract
|
|
import fitz
|
|
import os
|
|
from config.settings import MEDIA_ROOT
|
|
from watchdog.observers import Observer
|
|
from watchdog.events import FileSystemEventHandler
|
|
from documents.models import Document
|
|
import logging
|
|
import time
|
|
|
|
|
|
class PDFHandler(FileSystemEventHandler):
|
|
def __init__(self):
|
|
logging.basicConfig(level=logging.INFO,
|
|
format='%(asctime)s - %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S')
|
|
|
|
self.logger = logging.getLogger(__name__)
|
|
self.logger.info("Starting Document Watcher...")
|
|
|
|
def on_created(self, event):
|
|
if event.is_directory:
|
|
return None
|
|
|
|
if event.src_path.endswith('.pdf'):
|
|
self.logger.info(f"New PDF file detected: {event.src_path}")
|
|
self.process_pdf(event.src_path)
|
|
|
|
def process_pdf(self, file_path):
|
|
try:
|
|
filename = os.path.basename(file_path)
|
|
metadata = ""
|
|
document_type = ""
|
|
|
|
with fitz.open(file_path) as doc:
|
|
num_pages = len(doc)
|
|
|
|
for page_num in range(num_pages):
|
|
page = doc[page_num]
|
|
pix = page.get_pixmap(matrix=(1.2, 1.2))
|
|
|
|
# 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()
|
|
|
|
lines = text.split('\n')
|
|
|
|
for line in lines:
|
|
if line.strip():
|
|
document_type = line.strip().lower()
|
|
break
|
|
if not document_type or document_type not in Document.DOCUMENT_TYPE_CHOICES:
|
|
document_type = "other"
|
|
|
|
metadata += text
|
|
|
|
document, created = Document.objects.get_or_create(
|
|
name=filename,
|
|
defaults={
|
|
'number_pages': num_pages,
|
|
'ocr_metadata': metadata,
|
|
'document_type': document_type
|
|
}
|
|
)
|
|
|
|
if created:
|
|
self.logger.info(f"Document '{filename}' created successfully with type '{
|
|
document_type}'.")
|
|
|
|
else:
|
|
self.logger.info(f"Document '{filename}' already exists.")
|
|
|
|
os.remove(file_path)
|
|
except Exception as e:
|
|
self.logger.error(f"Error processing PDF: {str(e)}")
|
|
|
|
|
|
class PDFWatcher:
|
|
def __init__(self):
|
|
self.observer = Observer()
|
|
|
|
def run(self):
|
|
event_handler = PDFHandler()
|
|
watch_directory = f"{MEDIA_ROOT}/uploads"
|
|
|
|
self.observer.schedule(
|
|
event_handler, watch_directory, recursive=True)
|
|
self.observer.start()
|
|
|
|
try:
|
|
while True:
|
|
time.sleep(5)
|
|
except:
|
|
self.observer.stop()
|
|
|
|
self.observer.join()
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Runs a dedicated file watcher service"
|
|
|
|
def handle(self, *args, **options):
|
|
watcher = PDFWatcher()
|
|
watcher.run()
|