chandra2 / utils_chandra.py
Hyungseoky's picture
Upload 8 files
9773110 verified
Raw
History Blame Contribute Delete
25.4 kB
# utils_chandra.py - μœ ν‹Έλ¦¬ν‹° λͺ¨λ“ˆ (κΈ°μ‘΄ utils_test.py 기반, config 참쑰만 λ³€κ²½)
import asyncio
import subprocess
import tempfile
import fitz # PyMuPDF
import base64
import re
import shutil
import traceback
from pathlib import Path
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from typing import List, Dict, Tuple, Any, Optional
import json
import logging
from concurrent.futures import ThreadPoolExecutor
# === config μ°Έμ‘° λ³€κ²½ ===
from config_chandra import LIBREOFFICE_CONTAINERS, MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR, THREAD_POOL, PICTURE_CROP_MARGIN
logger = logging.getLogger(__name__)
# === μ „μ—­ μ»¨ν…Œμ΄λ„ˆ 건강도 좔적 ===
_container_failure_count = {name: 0 for name in LIBREOFFICE_CONTAINERS}
_container_lock = asyncio.Lock()
_last_libre_index = -1
async def get_healthy_container(base_name: str = None) -> str:
async with _container_lock:
global _last_libre_index
candidates = [
name for name in LIBREOFFICE_CONTAINERS
if _container_failure_count.get(name, 0) < 5
]
if not candidates:
return min(LIBREOFFICE_CONTAINERS, key=lambda x: _container_failure_count.get(x, 0))
if base_name in candidates:
next_idx = (candidates.index(base_name) + 1) % len(candidates)
else:
next_idx = (_last_libre_index + 1) % len(candidates)
_last_libre_index = next_idx
return candidates[_last_libre_index]
# === sorting μœ ν‹Έλ¦¬ν‹° ===
def remove_close_values(sorted_list, threshold=30):
if len(sorted_list) <= 1:
return sorted_list[:]
sorted_list = sorted(sorted_list)
result = [sorted_list[0]]
for i in range(1, len(sorted_list)):
current = sorted_list[i]
last_kept = result[-1]
if current - last_kept <= threshold:
continue
else:
result.append(current)
return result
def cluster_x_positions(x_positions: List[float], text_items: List[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
MIN_COLUMN_GAP = 200
MIN_COLUMN_BOUNDARY = int(float(min(item['bbox'][2] for item in text_items)) * 0.6)
MAX_COLUMN_BOUNDARY = int(float(max(item['bbox'][2] for item in text_items)) * 0.7)
THRESHOLD = 50
X_MARGIN = 40
if len(x_positions) <= 1:
return [text_items]
gaps = []
for i in range(len(x_positions) - 1):
gap = abs(x_positions[i + 1] - x_positions[i])
gaps.append((gap, max(x_positions[i], x_positions[i + 1])))
gaps.sort(reverse=True)
column_boundaries = []
for gap, boundary in gaps:
is_significant_gap = False
if len(set(gaps)) == 1:
is_significant_gap = gap >= MIN_COLUMN_GAP
else:
avg_gap = sum(g[0] for g in gaps) / len(gaps)
relative_condition = gap >= avg_gap
absolute_condition = (boundary >= MIN_COLUMN_BOUNDARY) & (boundary <= MAX_COLUMN_BOUNDARY) & (gap > MIN_COLUMN_GAP)
is_significant_gap = relative_condition and absolute_condition
if is_significant_gap:
column_boundaries.append(boundary)
column_boundaries = remove_close_values(column_boundaries, THRESHOLD)
column_boundaries = (np.array(column_boundaries) - X_MARGIN).tolist()
column_boundaries.sort()
return len(column_boundaries), column_boundaries
def analyze_document_layout(layout_data: List[Dict[str, Any]]) -> str:
column_boundaries = [1200]
text_items = [item for item in layout_data if item['category'] in ['Text', 'List-item']]
if len(text_items) < 2:
return ['single_column', column_boundaries]
sorted_text_items = sorted(text_items, key=lambda item: item['bbox'][1])
x_positions = [item['bbox'][0] for item in sorted_text_items]
num_of_column_boundaries, column_boundaries = cluster_x_positions(x_positions, text_items)
if num_of_column_boundaries == 0:
return ['single_column', column_boundaries]
elif num_of_column_boundaries == 1:
return ['two_column', column_boundaries]
elif num_of_column_boundaries >= 2:
return ['three_column', column_boundaries]
return ['single_column', column_boundaries]
def column_preparation(column_boundaries: List[int], text_items: List[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
columns = []
if not column_boundaries:
columns.append(text_items)
else:
for i in range(len(column_boundaries) + 1):
if i == 0:
column_items = [item for item in text_items if item['bbox'][0] < column_boundaries[0]]
elif i == len(column_boundaries):
column_items = [item for item in text_items if item['bbox'][0] >= column_boundaries[-1]]
else:
column_items = [item for item in text_items if column_boundaries[i - 1] <= item['bbox'][0] < column_boundaries[i]]
if column_items:
columns.append(column_items)
return columns
def get_page_width(layout_data: List[Dict[str, Any]]) -> float:
max_right = max(item['bbox'][2] for item in layout_data)
return max_right
def is_full_width_header(header: Dict[str, Any], layout_data: List[Dict[str, Any]]) -> bool:
page_width = get_page_width(layout_data)
header_width = header['bbox'][2] - header['bbox'][0]
header_left = header['bbox'][0]
width_ratio = header_width / page_width
is_wide = width_ratio > 0.7
is_left_aligned = header_left < page_width * 0.2
return is_wide and is_left_aligned
def sort_single_column(layout_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
return sorted(layout_data, key=lambda x: (x['bbox'][1], x['bbox'][0]))
def sort_multi_column(layout_data: List[Dict[str, Any]], column_boundaries: List[int]) -> List[Dict[str, Any]]:
all_headers = [item for item in layout_data if item['category'] in ['Page-header', 'Section-header']]
text_items = [item for item in layout_data if item['category'] not in ['Page-header', 'Section-header']]
full_width_headers = [h for h in all_headers if is_full_width_header(h, layout_data)]
partial_headers = [h for h in all_headers if not is_full_width_header(h, layout_data)]
all_text_items = text_items + partial_headers
sorted_section_headers = sorted(full_width_headers, key=lambda x: x['bbox'][1])
if not all_text_items:
return sorted_section_headers
columns = column_preparation(column_boundaries, all_text_items)
columns.sort(key=lambda col: min(item['bbox'][0] for item in col))
for column in columns:
column.sort(key=lambda x: x['bbox'][1])
result = []
if not sorted_section_headers:
for column in columns:
result.extend(column)
return result
first_header_top = sorted_section_headers[0]['bbox'][1]
for column in columns:
pre_header_items = [item for item in column if item['bbox'][1] < first_header_top]
result.extend(pre_header_items)
for item in pre_header_items:
if item in column:
column.remove(item)
for i, header in enumerate(sorted_section_headers):
result.append(header)
header_bottom = header['bbox'][3]
next_header_top = float('inf')
if i + 1 < len(sorted_section_headers):
next_header_top = sorted_section_headers[i + 1]['bbox'][1]
for column in columns:
section_items = [item for item in column if header_bottom <= item['bbox'][1] < next_header_top]
result.extend(section_items)
for item in section_items:
if item in column:
column.remove(item)
for column in columns:
result.extend(column)
return result
def sort_two_column(layout_data, column_boundaries):
return sort_multi_column(layout_data, column_boundaries)
def sort_three_column(layout_data, column_boundaries):
return sort_multi_column(layout_data, column_boundaries)
def smart_sort_layout(layout_data: List[Dict[str, Any]]) -> List[Dict[str, any]]:
if not layout_data:
return []
layout_inf = analyze_document_layout(layout_data)
layout_type, column_boundaries = layout_inf
if layout_type == 'single_column':
return sort_single_column(layout_data)
elif layout_type == 'two_column':
return sort_two_column(layout_data, column_boundaries)
elif layout_type == 'three_column':
return sort_three_column(layout_data, column_boundaries)
else:
return sort_single_column(layout_data)
# === 이미지 리사이징 μœ ν‹Έλ¦¬ν‹° ===
def round_by_factor(number: int, factor: int) -> int:
return round(number / factor) * factor
def smart_resize(height, width, factor=28, min_pixels=3136, max_pixels=11289600):
if max(height, width) / min(height, width) > 200:
raise ValueError("Aspect ratio must be < 200")
h_bar = max(factor, round_by_factor(height, factor))
w_bar = max(factor, round_by_factor(width, factor))
total_pixels = h_bar * w_bar
if total_pixels > max_pixels:
beta = (height * width / max_pixels) ** 0.5
h_bar = round_by_factor(height / beta, factor)
w_bar = round_by_factor(width / beta, factor)
elif total_pixels < min_pixels:
beta = (min_pixels / (height * width)) ** 0.5
h_bar = round_by_factor(height * beta, factor)
w_bar = round_by_factor(width * beta, factor)
return h_bar, w_bar
def fetch_image(image_input, min_pixels=None, max_pixels=None):
image = image_input.convert('RGB')
min_pixels = min_pixels or MIN_PIXELS
max_pixels = max_pixels or MAX_PIXELS
h, w = smart_resize(image.height, image.width, IMAGE_FACTOR, min_pixels, max_pixels)
resized = image.resize((w, h), Image.LANCZOS)
return resized
def load_images_from_pdf(pdf_path: str, highqual: bool = False) -> List[Image.Image]:
dpi_scale = 3.0 if highqual else 1.5
images = []
try:
pdf_document = fitz.open(pdf_path)
mat = fitz.Matrix(dpi_scale, dpi_scale)
for page_num in range(len(pdf_document)):
page = pdf_document.load_page(page_num)
pix = page.get_pixmap(matrix=mat)
img_data = pix.tobytes("ppm")
image = Image.open(BytesIO(img_data)).convert('RGB')
images.append(image)
pdf_document.close()
except Exception as e:
logger.error(f"PDF λ‘œλ“œ μ‹€νŒ¨: {e}")
raise
return images
def scale_bbox(bbox, from_size, to_size):
x1, y1, x2, y2 = bbox
curr_w, curr_h = from_size
orig_w, orig_h = to_size
x1 = int(x1 * orig_w / curr_w)
y1 = int(y1 * orig_h / curr_h)
x2 = int(x2 * orig_w / curr_w)
y2 = int(y2 * orig_h / curr_h)
return [x1, y1, x2, y2]
def draw_layout_on_image_with_scale(image, layout_data, resized_size):
colors = {
'Caption': '#FF6B6B', 'Footnote': '#4ECDC4', 'Formula': '#45B7D1',
'List-item': '#96CEB4', 'Page-footer': '#FFEAA7', 'Page-header': '#DDA0DD',
'Picture': '#FFD93D', 'Section-header': '#6C5CE7', 'Table': '#FD79A8',
'Text': '#74B9FF', 'Title': '#E17055'
}
orig_size = image.size
img_copy = image.copy()
draw = ImageDraw.Draw(img_copy)
margin = PICTURE_CROP_MARGIN
try:
try:
font = ImageFont.truetype("/usr/share/fonts/truetype/nanum/NanumGothic.ttf", 12)
except Exception:
font = ImageFont.load_default()
for item in layout_data:
if 'bbox' not in item or 'category' not in item:
continue
raw_bbox = item['bbox']
bbox = scale_bbox(raw_bbox, from_size=resized_size, to_size=orig_size)
category = item['category']
color = colors.get(category, '#000000')
x1, y1, x2, y2 = map(int, bbox)
if category == 'Picture' and margin > 0:
x1_expanded = max(0, x1 - margin)
y1_expanded = max(0, y1 - margin)
x2_expanded = min(orig_size[0], x2 + margin)
y2_expanded = min(orig_size[1], y2 + margin)
draw.rectangle([x1_expanded, y1_expanded, x2_expanded, y2_expanded], outline=color, width=2)
draw.rectangle(bbox, outline='#888888', width=1)
else:
draw.rectangle(bbox, outline=color, width=2)
label = category
label_bbox = draw.textbbox((0, 0), label, font=font)
label_width = label_bbox[2] - label_bbox[0]
label_height = label_bbox[3] - label_bbox[1]
label_x = x1
label_y = max(0, y1 - label_height - 2)
draw.rectangle(
[label_x, label_y, label_x + label_width + 4, label_y + label_height + 2],
fill=color
)
draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
except Exception as e:
logger.error(f"Error drawing layout: {e}")
return img_copy
def correct_markdown_syntax(md_text: str) -> str:
if not md_text or not isinstance(md_text, str):
return md_text
formula_blocks = []
def save_formula(match):
formula_blocks.append(match.group(0))
return f"__FORMULA_{len(formula_blocks)-1}__"
md_text = re.sub(r'\$\$.*?\$\$', save_formula, md_text, flags=re.DOTALL | re.MULTILINE)
md_text = re.sub(r'^#{5,}\s*(.+)$', r'##### \1', md_text, flags=re.MULTILINE)
md_text = re.sub(r'^(\s*)[-*]\s+[*-]\s+(.+)$', r'\1- \2', md_text, flags=re.MULTILINE)
md_text = re.sub(r'^(\s*)(#{1,5})\s+#{1,}\s+(.+)$', r'\1\2 \3', md_text, flags=re.MULTILINE)
md_text = re.sub(r'^(\s*)\*\*([^*\n].*)$', r'\1\2', md_text, flags=re.MULTILINE)
md_text = re.sub(r'^(\s*)\*([^*\n].*)$', r'\1\2', md_text, flags=re.MULTILINE)
md_text = re.sub(r'\n{3,}', '\n\n', md_text)
md_text = re.sub(r'[ \t]+$', '', md_text, flags=re.MULTILINE)
md_text = re.sub(r'^(\s*)- (\d+\. )', r'\1\2', md_text, flags=re.MULTILINE)
for i, block in enumerate(formula_blocks):
md_text = md_text.replace(f"__FORMULA_{i}__", block)
return md_text.strip()
def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
markdown_lines = []
sorted_items = smart_sort_layout(layout_data)
for item in sorted_items:
category = item.get('category', '')
text = item.get(text_key, '')
bbox = item.get('bbox', [])
if category == 'Picture':
if bbox and len(bbox) == 4:
x1, y1, x2, y2 = map(int, bbox)
x_margin = int(0.05 * float(abs(x2 - x1)))
y_margin = int(0.05 * float(abs(y2 - y1)))
x1, y1 = max(0, x1 - x_margin), max(0, y1 - y_margin)
x2, y2 = min(image.width, x2 + x_margin), min(image.height, y2 + y_margin)
if x2 > x1 and y2 > y1:
cropped_img = image.crop((x1, y1, x2, y2))
buffer = BytesIO()
cropped_img.save(buffer, format='PNG')
img_data = base64.b64encode(buffer.getvalue()).decode()
markdown_lines.append(f"![Image](data:image/png;base64,{img_data})\n")
else:
markdown_lines.append("![Image](Invalid region)\n")
else:
markdown_lines.append("![Image](Image detected)\n")
elif not text:
continue
elif category == 'Title':
markdown_lines.append(f"# {text}\n")
elif category == 'Section-header':
markdown_lines.append(f"## {text}\n")
elif category == 'Text':
markdown_lines.append(f"{text}\n")
elif category == 'List-item':
lines = text.strip().split('\n')
for line in lines:
line = line.strip()
if not line:
continue
if re.match(r'^(\d+\.|\*\s?|-\s?|[a-zA-Z]\.\s?|β€’)', line):
markdown_lines.append(f"- {line}")
else:
markdown_lines.append(f" {line}")
markdown_lines.append("")
elif category == 'Table':
if text.strip().startswith('<'):
markdown_lines.append(f"{text}\n")
else:
markdown_lines.append(f"**Table:** {text}\n")
elif category == 'Formula':
if text.strip().startswith('$') or '\\' in text:
markdown_lines.append(f"$$\n{text}\n$$\n")
else:
markdown_lines.append(f"**Formula:** {text}\n")
elif category == 'Caption':
markdown_lines.append(f"*{text}*\n")
elif category == 'Footnote':
markdown_lines.append(f"^{text}^\n")
elif category in ['Page-header', 'Page-footer']:
continue
else:
markdown_lines.append(f"{text}\n")
markdown_lines.append("")
return "\n".join(markdown_lines)
def PILimage_to_base64(image: Image.Image, format: str = 'PNG') -> str:
if isinstance(image, str):
image = Image.open(image)
buffered = BytesIO()
image.save(buffered, format=format, optimize=True)
base64_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
return f"data:image/{format.lower()};base64,{base64_str.strip()}"
# === 비동기 λ³€ν™˜ ν•¨μˆ˜ ===
LIBREOFFICE_SEMAPHORE = asyncio.Semaphore(3)
async def async_convert_office_to_pdf(
input_path: str,
container_name: str = None,
max_retries: int = 3,
temp_dir: Optional[Path] = None
) -> str:
if container_name is None:
container_name = LIBREOFFICE_CONTAINERS[0]
elif container_name not in LIBREOFFICE_CONTAINERS:
logger.warning(f"잘λͺ»λœ μ»¨ν…Œμ΄λ„ˆ 이름: {container_name}, κΈ°λ³Έκ°’ μ‚¬μš©")
container_name = LIBREOFFICE_CONTAINERS[0]
path_obj = Path(input_path)
if not path_obj.exists():
raise FileNotFoundError(f"μž…λ ₯ 파일 μ—†μŒ: {input_path}")
own_temp_dir = False
if temp_dir is None:
temp_dir = Path(tempfile.mkdtemp(prefix="office2pdf_"))
own_temp_dir = True
else:
temp_dir = Path(temp_dir)
temp_dir.mkdir(exist_ok=True)
output_pdf_name = f"{path_obj.stem}.pdf"
host_pdf_path = temp_dir / output_pdf_name
container_input_path = f"/config/{path_obj.name}"
container_output_path = f"/config/{output_pdf_name}"
for attempt in range(1, max_retries + 1):
current_container = container_name if attempt == 1 else await get_healthy_container(container_name)
timeout_factor = 1.5 ** (attempt - 1)
copy_timeout = 80.0 * timeout_factor
convert_timeout = 400.0 * timeout_factor
try:
async with LIBREOFFICE_SEMAPHORE:
logger.info(f"[Try {attempt}/{max_retries}] λ³€ν™˜: {path_obj.name} β†’ {current_container}")
proc = await asyncio.create_subprocess_exec(
"docker", "cp", str(path_obj), f"{current_container}:{container_input_path}",
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
)
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=copy_timeout)
except asyncio.TimeoutError:
raise RuntimeError("docker cp timeout")
if proc.returncode != 0:
stderr_msg = stderr.decode().strip() if stderr else "Unknown error"
raise RuntimeError(f"Copy to container failed: {stderr_msg}")
proc = await asyncio.create_subprocess_exec(
"docker", "exec", current_container,
"soffice", "--headless", "--convert-to", "pdf", "--outdir", "/config",
container_input_path,
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
)
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=convert_timeout)
except asyncio.TimeoutError:
raise RuntimeError("Conversion timeout")
if proc.returncode != 0:
stderr_msg = stderr.decode().strip() if stderr else "Unknown error"
raise RuntimeError(f"Conversion failed: {stderr_msg}")
proc = await asyncio.create_subprocess_exec(
"docker", "exec", current_container, "test", "-f", container_output_path
)
await proc.wait()
if proc.returncode != 0:
proc = await asyncio.create_subprocess_exec(
"docker", "exec", current_container,
"find", "/tmp", "-name", f"{path_obj.stem}*.pdf", "-type", "f",
stdout=asyncio.subprocess.PIPE
)
stdout, _ = await proc.communicate()
found_paths = stdout.decode().strip().splitlines()
if not found_paths:
raise FileNotFoundError("PDF νŒŒμΌμ„ 찾을 수 μ—†μŒ")
container_output_path = found_paths[0]
proc = await asyncio.create_subprocess_exec(
"docker", "cp", f"{current_container}:{container_output_path}", str(host_pdf_path),
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
)
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=copy_timeout)
except asyncio.TimeoutError:
raise RuntimeError("docker cp from container timeout")
if proc.returncode != 0 or not host_pdf_path.exists():
raise FileNotFoundError(f"κ²°κ³Ό 파일 μ—†μŒ: {host_pdf_path}")
async with _container_lock:
_container_failure_count[current_container] = max(0, _container_failure_count[current_container] - 1)
logger.info(f"λ³€ν™˜ 성곡: {host_pdf_path}")
return str(host_pdf_path)
except Exception as e:
logger.warning(f"λ³€ν™˜ μ‹€νŒ¨ (μ‹œλ„ {attempt}, μ»¨ν…Œμ΄λ„ˆ: {current_container}): {e}")
async with _container_lock:
_container_failure_count[current_container] += 1
if attempt == max_retries:
if own_temp_dir:
try:
shutil.rmtree(temp_dir, ignore_errors=True)
except:
pass
raise
await asyncio.sleep(2 ** attempt)
raise RuntimeError("λͺ¨λ“  μž¬μ‹œλ„ μ‹€νŒ¨")
# === 비동기 래퍼 ===
async def async_fetch_image(image, min_pixels=None, max_pixels=None):
return await asyncio.get_event_loop().run_in_executor(
THREAD_POOL, fetch_image, image, min_pixels, max_pixels
)
async def async_load_images_from_pdf(pdf_path, highqual=False):
return await asyncio.get_event_loop().run_in_executor(
THREAD_POOL, load_images_from_pdf, pdf_path, highqual
)
async def async_draw_layout_on_image_with_scale(image, layout_data, resized_size):
return await asyncio.get_event_loop().run_in_executor(
THREAD_POOL, draw_layout_on_image_with_scale, image, layout_data, resized_size
)
async def async_correct_markdown_syntax(md_text):
return await asyncio.get_event_loop().run_in_executor(
THREAD_POOL, correct_markdown_syntax, md_text
)
async def async_layoutjson2md(image, layout_data, text_key='text'):
return await asyncio.get_event_loop().run_in_executor(
THREAD_POOL, layoutjson2md, image, layout_data, text_key
)
# === 톡합 λ‘œλ“œ ν•¨μˆ˜ ===
async def async_load_file_for_preview(
file_path: str,
container_name: str = "libreoffice",
highqual: bool = False,
temp_dir: Optional[Path] = None
) -> Tuple[Optional[Image.Image], str, Optional[str]]:
path_obj = Path(file_path)
ext = path_obj.suffix.lower()
try:
if ext in ['.doc', '.docx', '.ppt', '.pptx', '.odt', '.ods', '.odp']:
logger.info(f"Office 파일 비동기 λ³€ν™˜ 쀑: {path_obj.name}")
selected_container = await get_healthy_container(container_name)
temp_pdf_path = await async_convert_office_to_pdf(
str(path_obj), container_name=selected_container, temp_dir=temp_dir
)
images = await async_load_images_from_pdf(temp_pdf_path, highqual=highqual)
if not images:
return None, "PDF λ‘œλ“œ μ‹€νŒ¨", temp_pdf_path
return images[0], f"Page 1 / {len(images)}", temp_pdf_path
elif ext == '.pdf':
images = await async_load_images_from_pdf(str(path_obj), highqual=highqual)
if not images:
return None, "PDF λ‘œλ“œ μ‹€νŒ¨", None
return images[0], f"Page 1 / {len(images)}", None
elif ext in ['.jpg', '.jpeg', '.png', '.bmp']:
image = Image.open(path_obj).convert('RGB')
return image, "Page 1 / 1", None
else:
return None, f"μ§€μ›ν•˜μ§€ μ•ŠλŠ” ν˜•μ‹: {ext}", None
except Exception as e:
logger.error(f"파일 λ‘œλ“œ 였λ₯˜: {e}")
traceback.print_exc()
return None, f"λ‘œλ“œ μ‹€νŒ¨: {str(e)}", None