def convert(input_filename, doc_id, slug):
# Provision a temporary directory in which to handle document conversion
document_directory = tempfile.mkdtemp(prefix=TMP_DIR)
# Grab file from storage to tmp
tmp_path = os.path.join(document_directory, Path(input_filename).name)
with storage.open(input_filename, "rb") as document_file:
with open(tmp_path, "wb") as tmp_file:
tmp_file.write(document_file.read())
# Run LibreOffice
libre_office_convert(tmp_path)
# Remove created file (early, just to free RAM that might be needed later)
os.remove(tmp_path)
# Put converted file back in storage
# We expect the filename to be the same but with a pdf extension
# (LibreOffice offers no mechanism to specify a precise name)
output_path = str(Path(tmp_path).with_suffix(".pdf"))
output_filename = path.doc_path(doc_id, slug)
with storage.open(output_filename, "wb") as output_document_file:
with open(output_path, "rb") as pdf_file:
output_document_file.write(pdf_file.read())
# Remove temporary directory
shutil.rmtree(document_directory)
def doc_embedding(project_id, language, tfidf, features, doc_svd):
"""Calculate the doc embeddings"""
import fasttext
logger.info("[SIDEKICK PREPROCESS] project_id: %s - doc embeddings",
project_id)
# Load the embedding model
# error if language not present
language = LANGUAGES.get(language, language)
model = fasttext.load_model(
os.path.join(EMBEDDING_DIR, f"cc.{language}.300.bin"))
embedding_vectors = np.array(
[model.get_word_vector(feature) for feature in features])
# scale embedding vectors based on frequency of the words
doc_embeddings = np.dot(tfidf.A, embedding_vectors)
# Doc vectors are just doc_svd and doc_embeddings concatenated
doc_vectors = np.hstack((doc_svd, doc_embeddings))
# Serialize doc vectors to file
with storage.open(path.sidekick_document_vectors_path(project_id),
"wb") as vectors_file:
np.savez_compressed(vectors_file, doc_vectors)
def get_document_vectors(self):
"""Fetch the pre-preocessed document vectors from storage"""
with storage.open(path.sidekick_document_vectors_path(self.project_id),
"rb") as vectors_file:
doc_vector_obj = np.load(vectors_file)
# Grab document vector matrix
return (doc_vector_obj.get("vectors"), doc_vector_obj.get("ids"))
def doc_embedding_(project_id, _language, _tfidf, _features, doc_svd, doc_ids):
"""Simpler doc embeddings - skip word vectors and just use the doc svd"""
logger.info("[SIDEKICK PREPROCESS] project_id: %s - doc embeddings",
project_id)
# Serialize doc vectors to file
with storage.open(path.sidekick_document_vectors_path(project_id),
"wb") as vectors_file:
np.savez_compressed(vectors_file, vectors=doc_svd, ids=doc_ids)
def get_text(self):
try:
return (storage.open(path.text_path(self.pk, self.slug),
"rb").read().decode("utf8"))
except ValueError as exc:
logger.error(
"Error getting text: Document: %d Exception: %s",
self.pk,
exc,
exc_info=sys.exc_info(),
)
return ""
def get_all_page_text(self):
try:
return json.loads(
storage.open(path.json_text_path(self.pk, self.slug),
"rb").read().decode("utf8"))
except ValueError as exc:
logger.error(
"Error getting all page text: Document: %d Exception: %s",
self.pk,
exc,
exc_info=sys.exc_info(),
)
return {"pages": [], "updated": None}
def handle(self, *args, **options):
data_files = os.listdir(TESSERACT_DATA_DIRECTORY)
print("UPLOADING", data_files)
print("...")
for data_file_path in data_files:
with open(os.path.join(TESSERACT_DATA_DIRECTORY, data_file_path),
"rb") as data_file:
with storage.open(
os.path.join(MINIO_DATA_DIRECTORY, data_file_path),
"wb") as minio_file:
minio_file.write(data_file.read())
print("WROTE ALL FILES")
def download_tmp_file(relative_path):
"""Downloads the requested data file to a tmp directory."""
Path(TMP_DIRECTORY).mkdir(
parents=True, exist_ok=True
) # Make tmp directory if it doesn't exist
local_file_path = os.path.join(TMP_DIRECTORY, relative_path)
if os.path.exists(local_file_path):
# OCR language pack already downloaded
return
# Check if tmp directory is too big
if local_folder_size(TMP_DIRECTORY) > TMP_SIZE_LIMIT:
# If so, just delete all OCR data (shouldn't happen too often)
logger.warning("[Deleting tmp OCR data]")
files = Path(TMP_DIRECTORY).rglob("*")
for file in files:
os.remove(file)
# Download OCR data file
with storage.open(
os.path.join(OCR_DATA_DIRECTORY, relative_path), "rb"
) as ocr_data_file, open(local_file_path, "wb") as local_file:
local_file.write(ocr_data_file.read())
def ocr_page(doc_id, page_path, upload_text_path, access, ocr_code="eng"):
"""Internal method to run OCR on a single page.
Returns:
The page text.
"""
# Download the requisite language data
logger.info("[OCR PAGE] doc_id %s", doc_id)
download_language_pack(ocr_code)
download_tmp_file(PDF_FONT_FILE)
logger.info("[OCR PAGE] download complete doc_id %s", doc_id)
# Initialize temporary files
tmp_files = {
"img": tempfile.mkstemp(suffix=".png")[1],
"pdf": tempfile.mkstemp()[1],
"text": tempfile.mkstemp()[1],
}
# Capture the page image as a temporary PNG file
with storage.open(page_path, "rb") as image_file:
img = Image.open(image_file).convert("RGB")
# Resize only if image is too big (OCR computation is slow with large images)
if img.width > DESIRED_WIDTH:
resize = DESIRED_WIDTH / img.width
img = img.resize(
(DESIRED_WIDTH, round(img.height * resize)), Image.ANTIALIAS
)
img.save(tmp_files["img"], "png")
logger.info("[OCR PAGE] image resized doc_id %s", doc_id)
# Use Tesseract OCR to render a text-only PDF and txt file
tess = Tesseract(ocr_code)
text = ""
pdf_contents = b""
try:
tess.create_renderer(tmp_files["pdf"], tmp_files["text"])
tess.render(tmp_files["img"])
tess.destroy_renderer()
logger.info("[OCR PAGE] rendered doc_id %s", doc_id)
# Get txt and text-only pdf file contents
with open(tmp_files["pdf"] + ".pdf", "rb") as pdf_file:
pdf_contents = pdf_file.read()
with storage.open(upload_text_path, "w", access=access) as new_text_file:
with open(tmp_files["text"] + ".txt", "r", encoding="utf-8") as text_file:
# Store text locally to return (gets used by Redis later)
text = text_file.read()
# Also upload text file to s3
new_text_file.write(text)
logger.info("[OCR PAGE] data stored doc_id %s", doc_id)
finally:
logger.info("[OCR PAGE] cleanup doc_id %s", doc_id)
os.remove(tmp_files["pdf"])
os.remove(tmp_files["text"])
os.remove(tmp_files["img"])
return text, pdf_contents