Spaces:
Sleeping
Sleeping
File size: 3,305 Bytes
c1c25ed |
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 |
## PDF processing up to text extraction
import os
import shutil
import fitz
from PIL import Image
import numpy as np
import cv2
import pytesseract
from pytesseract import Output
import zipfile
from pdf2image import convert_from_path
import json
def convert_to_rgb(image_path):
img = Image.open(image_path)
rgb_img = img.convert("RGB")
return rgb_img
def preprocess_image(image):
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
denoised = cv2.fastNlMeansDenoising(binary, None, 30, 7, 21)
resized = cv2.resize(denoised, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
return resized
def extract_vertical_blocks(image):
image_np = np.array(image)
data = pytesseract.image_to_data(image_np, lang='fra', output_type=Output.DICT)
blocks = []
current_block = ""
current_block_coords = [float('inf'), float('inf'), 0, 0]
last_bottom = -1
line_height = 0
for i in range(len(data['text'])):
if int(data['conf'][i]) > 0:
text = data['text'][i]
x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
if line_height == 0:
line_height = h * 1.2
if y > last_bottom + line_height:
if current_block:
blocks.append({
"text": current_block.strip(),
"coords": current_block_coords
})
current_block = ""
current_block_coords = [float('inf'), float('inf'), 0, 0]
current_block += text + " "
current_block_coords[0] = min(current_block_coords[0], x)
current_block_coords[1] = min(current_block_coords[1], y)
current_block_coords[2] = max(current_block_coords[2], x + w)
current_block_coords[3] = max(current_block_coords[3], y + h)
last_bottom = y + h
if current_block:
blocks.append({
"text": current_block.strip(),
"coords": current_block_coords
})
return blocks
def draw_blocks_on_image(image_path, blocks, output_path):
image = cv2.imread(image_path)
for block in blocks:
coords = block['coords']
cv2.rectangle(image, (coords[0], coords[1]), (coords[2], coords[3]), (0, 0, 255), 2)
cv2.imwrite(output_path, image)
return output_path
def process_image(image, output_folder, page_number):
image = convert_to_rgb(image)
blocks = extract_vertical_blocks(image)
base_name = f'page_{page_number + 1}.png'
image_path = os.path.join(output_folder, base_name)
image.save(image_path)
annotated_image_path = os.path.join(output_folder, f'annotated_{base_name}')
annotated_image_path = draw_blocks_on_image(image_path, blocks, annotated_image_path)
return blocks, annotated_image_path
def save_extracted_text(blocks, page_number, output_folder):
text_file_path = os.path.join(output_folder, 'extracted_text.txt')
with open(text_file_path, 'a', encoding='utf-8') as f:
f.write(f"[PAGE {page_number}]\n")
for block in blocks:
f.write(block['text'] + "\n")
f.write("[FIN DE PAGE]\n\n")
return text_file_path |