danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Big code refactoring

Browse source
This commit is contained in:
pjmalandrino 2025-05-30 14:37:56 +02:00
parent 224e97ccc4
commit ffa2f168e1
8 changed files with 1231 additions and 2869 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1588
backend/app.py
View file

File diff suppressed because it is too large Load diff

318
backend/batch_treatment/batch_processor.py
View file

@ -1,7 +1,6 @@
#!/usr/bin/env python3
"""
Batch Processor for DocTags
Handles batch processing of PDF documents with parallel processing support
Batch Processor for DocTags - Handles parallel processing of PDF documents
"""
import os
@ -14,18 +13,15 @@ import logging
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import subprocess
import shutil
import zipfile
# Add the parent directory to Python path to allow imports
# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
from backend.utils import ensure_results_folder, run_command_with_timeout, format_duration
from backend.config import BATCH_WORKERS, PROCESSING_TIMEOUT
logger = logging.getLogger(__name__)
@ -47,7 +43,7 @@ class BatchProcessor:
'completed': False,
'paused': False,
'cancelled': False,
'page_statuses': {},
'page_statuses': {str(page): 'pending' for page in range(start_page, end_page + 1)},
'stages': {
'analysis': {'completed': 0, 'total': self.total_pages},
'visualization': {'completed': 0, 'total': self.total_pages},
@ -62,17 +58,13 @@ class BatchProcessor:
'logs': []
}
# Initialize page statuses
for page in range(start_page, end_page + 1):
self.state['page_statuses'][str(page)] = 'pending'
# Threading
self.lock = threading.Lock()
self.pause_event = threading.Event()
self.pause_event.set() # Start unpaused
# Create batch results directory
self.results_dir = Path("results") / f"batch_{batch_id}"
self.results_dir = ensure_results_folder() / f"batch_{batch_id}"
self.results_dir.mkdir(parents=True, exist_ok=True)
# Log file
@ -89,11 +81,11 @@ class BatchProcessor:
with self.lock:
self.state['logs'].append(log_entry)
# Keep only last 100 log entries in memory
# Keep only last 100 log entries
if len(self.state['logs']) > 100:
self.state['logs'] = self.state['logs'][-100:]
# Also write to log file
# Write to log file
with open(self.log_file, 'a') as f:
f.write(f"[{timestamp}] [{level.upper()}] {message}\n")
@ -125,7 +117,7 @@ class BatchProcessor:
raise Exception("Analyzer failed")
self.update_stage_progress('analysis')
# Check if paused or cancelled
# Check pause/cancel
self.pause_event.wait()
if self.state['cancelled']:
return False
@ -135,7 +127,7 @@ class BatchProcessor:
raise Exception("Visualizer failed")
self.update_stage_progress('visualization')
# Check if paused or cancelled
# Check pause/cancel
self.pause_event.wait()
if self.state['cancelled']:
return False
@ -152,7 +144,6 @@ class BatchProcessor:
self.update_page_status(page_num, 'completed')
self.log_message(f"Successfully processed page {page_num}", 'success')
return True
except Exception as e:
@ -167,45 +158,26 @@ class BatchProcessor:
self.update_page_status(page_num, 'failed')
self.log_message(f"Failed to process page {page_num}: {str(e)}", 'error')
return False
def run_analyzer(self, page_num):
"""Run the analyzer for a specific page"""
try:
# Use standard results directory for analyzer output
output_base = Path("results") / f"output"
# Update path to use backend directory
command = [
"python", "backend/page_treatment/analyzer.py",
"--image", self.pdf_file,
"--page", str(page_num),
"--output", str(output_base),
"--start-page", str(page_num),
"--end-page", str(page_num)
]
command = (f"python backend/page_treatment/analyzer.py "
f"--image {self.pdf_file} --page {page_num} "
f"--start-page {page_num} --end-page {page_num}")
self.log_message(f"Running analyzer for page {page_num}")
# Run the command
process = subprocess.Popen(
command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True
)
success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
# Send "n" to bypass prompts
stdout, stderr = process.communicate(input="n\n", timeout=300) # 5 minute timeout
if process.returncode != 0:
if not success:
raise Exception(f"Analyzer failed: {stderr}")
# Copy doctags to batch directory for archiving
doctags_src = Path("results") / "output.doctags.txt"
# Copy doctags to batch directory
doctags_src = ensure_results_folder() / "output.doctags.txt"
doctags_dst = self.results_dir / f"page_{page_num}.doctags.txt"
if doctags_src.exists():
shutil.copy2(doctags_src, doctags_dst)
self.log_message(f"DocTags saved for page {page_num}")
@ -221,47 +193,33 @@ class BatchProcessor:
def run_visualizer(self, page_num):
"""Run the visualizer for a specific page"""
try:
# The visualizer expects doctags in the standard location
doctags_path = Path("results") / "output.doctags.txt"
# First, ensure we have the right doctags file for this page
# Ensure correct doctags file is in place
doctags_path = ensure_results_folder() / "output.doctags.txt"
page_doctags = self.results_dir / f"page_{page_num}.doctags.txt"
if page_doctags.exists():
shutil.copy2(page_doctags, doctags_path)
# Update path to use backend directory
command = [
"python", "backend/page_treatment/visualizer.py",
"--doctags", str(doctags_path),
"--pdf", self.pdf_file,
"--page", str(page_num)
]
command = (f"python backend/page_treatment/visualizer.py "
f"--doctags {doctags_path} --pdf {self.pdf_file} --page {page_num}")
if self.options.get('adjust', True):
command.append("--adjust")
command += " --adjust"
self.log_message(f"Running visualizer for page {page_num}")
# Run the command
process = subprocess.run(
command,
capture_output=True,
text=True,
timeout=300
)
success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
if process.returncode != 0:
raise Exception(f"Visualizer failed: {process.stderr}")
if not success:
raise Exception(f"Visualizer failed: {stderr}")
# The visualizer should have created the file in results/
viz_src = Path("results") / f"visualization_page_{page_num}.png"
# Copy visualization to batch directory
viz_src = ensure_results_folder() / f"visualization_page_{page_num}.png"
viz_dst = self.results_dir / f"visualization_page_{page_num}.png"
if viz_src.exists():
shutil.copy2(viz_src, viz_dst)
self.log_message(f"Visualization saved for page {page_num}")
else:
self.log_message(f"Warning: Visualization file not found for page {page_num}", 'warning')
return True
@ -272,39 +230,29 @@ class BatchProcessor:
def run_extractor(self, page_num):
"""Run the picture extractor for a specific page"""
try:
# Ensure we have the right doctags file for this page
doctags_path = Path("results") / "output.doctags.txt"
# Ensure correct doctags file is in place
doctags_path = ensure_results_folder() / "output.doctags.txt"
page_doctags = self.results_dir / f"page_{page_num}.doctags.txt"
if page_doctags.exists():
shutil.copy2(page_doctags, doctags_path)
# Update path to use backend directory
command = [
"python", "backend/page_treatment/picture_extractor.py",
"--doctags", str(doctags_path),
"--pdf", self.pdf_file,
"--page", str(page_num)
]
command = (f"python backend/page_treatment/picture_extractor.py "
f"--doctags {doctags_path} --pdf {self.pdf_file} --page {page_num}")
if self.options.get('adjust', True):
command.append("--adjust")
command += " --adjust"
self.log_message(f"Running extractor for page {page_num}")
# Run the command
process = subprocess.run(
command,
capture_output=True,
text=True,
timeout=300
)
success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
if process.returncode != 0:
self.log_message(f"Extractor warning for page {page_num}: {process.stderr}", 'warning')
if not success:
self.log_message(f"Extractor warning for page {page_num}: {stderr}", 'warning')
# Count and copy extracted images
image_count = 0
pics_src = Path("results") / "pictures"
pics_src = ensure_results_folder() / "pictures"
pics_dst = self.results_dir / f"pictures_page_{page_num}"
if pics_src.exists():
@ -313,8 +261,8 @@ class BatchProcessor:
shutil.rmtree(pics_dst)
shutil.copytree(pics_src, pics_dst)
# Also copy to page-specific location for web interface
pics_web = Path("results") / f"pictures_page_{page_num}"
# Also copy for web interface
pics_web = ensure_results_folder() / f"pictures_page_{page_num}"
if pics_web.exists():
shutil.rmtree(pics_web)
shutil.copytree(pics_src, pics_web)
@ -322,8 +270,6 @@ class BatchProcessor:
# Count PNG files
image_count = len(list(pics_dst.glob("*.png")))
self.log_message(f"Extracted {image_count} images from page {page_num}")
else:
self.log_message(f"No images extracted from page {page_num}", 'info')
return image_count
@ -335,10 +281,11 @@ class BatchProcessor:
"""Main batch processing loop"""
try:
self.state['status'] = 'processing'
self.log_message(f"Starting batch processing for {self.pdf_file} (pages {self.start_page}-{self.end_page})")
self.log_message(f"Starting batch processing for {self.pdf_file} "
f"(pages {self.start_page}-{self.end_page})")
# Determine number of workers
max_workers = 4 if self.options.get('parallel', True) else 1
max_workers = BATCH_WORKERS if self.options.get('parallel', True) else 1
# Create page list
pages = list(range(self.start_page, self.end_page + 1))
@ -347,7 +294,8 @@ class BatchProcessor:
if max_workers > 1:
# Parallel processing
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {executor.submit(self.process_page, page): page for page in pages}
futures = {executor.submit(self.process_page, page): page
for page in pages}
for future in as_completed(futures):
if self.state['cancelled']:
@ -358,7 +306,8 @@ class BatchProcessor:
try:
future.result()
except Exception as e:
self.log_message(f"Unexpected error processing page {page}: {str(e)}", 'error')
self.log_message(f"Unexpected error processing page {page}: {str(e)}",
'error')
else:
# Sequential processing
for page in pages:
@ -376,7 +325,8 @@ class BatchProcessor:
self.state['status'] = 'cancelled' if self.state['cancelled'] else 'completed'
duration = time.time() - self.state['start_time']
self.log_message(f"Batch processing completed in {self.format_duration(duration)}", 'success')
self.log_message(f"Batch processing completed in {format_duration(duration)}",
'success')
except Exception as e:
self.log_message(f"Critical error in batch processing: {str(e)}", 'error')
@ -385,157 +335,100 @@ class BatchProcessor:
self.state['status'] = 'error'
def generate_report(self):
"""Generate a comprehensive HTML report"""
"""Generate HTML report of batch processing results"""
try:
self.log_message("Generating batch processing report")
report_path = self.results_dir / "report.html"
# Calculate statistics
duration = time.time() - self.state['start_time']
success_rate = (self.state['results']['successful'] / self.total_pages * 100) if self.total_pages > 0 else 0
success_rate = (self.state['results']['successful'] / self.total_pages * 100
if self.total_pages > 0 else 0)
html = f"""<!DOCTYPE html>
# Create report HTML
report_html = self._create_report_html(duration, success_rate)
# Save report
report_path = self.results_dir / "report.html"
with open(report_path, 'w') as f:
f.write(report_html)
self.log_message("Report generated successfully")
except Exception as e:
self.log_message(f"Error generating report: {str(e)}", 'error')
def _create_report_html(self, duration, success_rate):
"""Create HTML content for the report"""
html = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<title>Batch Processing Report - {self.pdf_file}</title>
<style>
body {{ font-family: Arial, sans-serif; margin: 40px; background: #f5f5f5; }}
.container {{ max-width: 1200px; margin: 0 auto; background: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }}
h1 {{ color: #2c3e50; margin-bottom: 10px; }}
.subtitle {{ color: #7f8c8d; margin-bottom: 30px; }}
.stats {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 20px; margin: 30px 0; }}
.container {{ max-width: 1200px; margin: 0 auto; background: white;
padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }}
h1 {{ color: #2c3e50; }}
.stats {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
gap: 20px; margin: 30px 0; }}
.stat-box {{ background: #ecf0f1; padding: 20px; border-radius: 8px; text-align: center; }}
.stat-value {{ font-size: 2.5em; font-weight: bold; color: #2c3e50; }}
.stat-label {{ color: #7f8c8d; margin-top: 5px; }}
.success {{ color: #27ae60; }}
.error {{ color: #e74c3c; }}
table {{ width: 100%; border-collapse: collapse; margin: 20px 0; }}
th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ecf0f1; }}
th {{ background: #34495e; color: white; }}
tr:hover {{ background: #f8f9fa; }}
.page-preview {{ display: inline-block; margin: 10px; text-align: center; }}
.page-preview img {{ max-width: 200px; max-height: 200px; border: 1px solid #ddd; }}
.failed-section {{ background: #fee; padding: 20px; border-radius: 8px; margin: 20px 0; }}
</style>
</head>
<body>
<div class="container">
<h1>Batch Processing Report</h1>
<p class="subtitle">Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
<p>Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
<div class="stats">
<div class="stat-box">
<div class="stat-value">{self.total_pages}</div>
<div class="stat-label">Total Pages</div>
<div>Total Pages</div>
</div>
<div class="stat-box">
<div class="stat-value success">{self.state['results']['successful']}</div>
<div class="stat-label">Successful</div>
<div>Successful</div>
</div>
<div class="stat-box">
<div class="stat-value error">{self.state['results']['failed']}</div>
<div class="stat-label">Failed</div>
<div>Failed</div>
</div>
<div class="stat-box">
<div class="stat-value">{self.state['results']['totalImages']}</div>
<div class="stat-label">Images Extracted</div>
<div>Images Extracted</div>
</div>
<div class="stat-box">
<div class="stat-value">{self.format_duration(duration)}</div>
<div class="stat-label">Processing Time</div>
<div class="stat-value">{format_duration(duration)}</div>
<div>Processing Time</div>
</div>
<div class="stat-box">
<div class="stat-value">{success_rate:.1f}%</div>
<div class="stat-label">Success Rate</div>
<div>Success Rate</div>
</div>
</div>
<h2>Processing Details</h2>
"""
# Add failed pages if any
if self.state['results']['failedPages']:
html += """
<h2>Failed Pages</h2>
<table>
<tr>
<th>Parameter</th>
<th>Value</th>
</tr>
<tr>
<td>PDF File</td>
<td>{self.pdf_file}</td>
</tr>
<tr>
<td>Page Range</td>
<td>{self.start_page} - {self.end_page}</td>
</tr>
<tr>
<td>Batch ID</td>
<td>{self.batch_id}</td>
</tr>
<tr>
<td>Parallel Processing</td>
<td>{'Enabled' if self.options.get('parallel', True) else 'Disabled'}</td>
</tr>
<tr>
<td>Auto-adjust Coordinates</td>
<td>{'Enabled' if self.options.get('adjust', True) else 'Disabled'}</td>
</tr>
</table>
<tr><th>Page Number</th><th>Reason</th></tr>
"""
for failed in self.state['results']['failedPages']:
html += f"<tr><td>{failed['pageNum']}</td><td>{failed['reason']}</td></tr>\n"
html += "</table>\n"
# Add failed pages section if any
if self.state['results']['failedPages']:
html += """
<div class="failed-section">
<h2>Failed Pages</h2>
<table>
<tr>
<th>Page Number</th>
<th>Reason</th>
</tr>
"""
for failed in self.state['results']['failedPages']:
html += f"""
<tr>
<td>{failed['pageNum']}</td>
<td>{failed['reason']}</td>
</tr>
"""
html += """
</table>
</div>
"""
# Add successful pages preview
html += """
<h2>Processed Pages</h2>
<div style="display: flex; flex-wrap: wrap; gap: 20px;">
"""
for page in range(self.start_page, self.end_page + 1):
if self.state['page_statuses'].get(str(page)) == 'completed':
viz_path = f"visualization_page_{page}.png"
html += f"""
<div class="page-preview">
<a href="{viz_path}" target="_blank">
<img src="{viz_path}" alt="Page {page}">
</a>
<p>Page {page}</p>
</div>
"""
html += """
</div>
html += """
</div>
</body>
</html>
"""
with open(report_path, 'w') as f:
f.write(html)
self.log_message("Report generated successfully")
except Exception as e:
self.log_message(f"Error generating report: {str(e)}", 'error')
return html
def pause(self):
"""Pause the batch processing"""
@ -552,11 +445,11 @@ class BatchProcessor:
def cancel(self):
"""Cancel the batch processing"""
self.state['cancelled'] = True
self.pause_event.set() # Ensure not stuck on pause
self.pause_event.set()
self.log_message("Batch processing cancelled")
def get_state(self):
"""Get the current state with calculated fields"""
"""Get current state with calculated fields"""
with self.lock:
state = self.state.copy()
@ -572,24 +465,12 @@ class BatchProcessor:
return state
def format_duration(self, seconds):
"""Format duration in seconds to human readable format"""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
if hours > 0:
return f"{hours}:{minutes:02d}:{secs:02d}"
else:
return f"{minutes}:{secs:02d}"
def create_zip_archive(self):
"""Create a ZIP archive of all results"""
"""Create ZIP archive of all results"""
try:
zip_path = self.results_dir / f"batch_results_{self.batch_id}.zip"
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
# Add all files in the results directory
for file_path in self.results_dir.rglob('*'):
if file_path.is_file() and file_path != zip_path:
arcname = file_path.relative_to(self.results_dir)
@ -648,4 +529,5 @@ def cleanup_old_batches(max_age_hours=24):
to_remove.append(batch_id)
for batch_id in to_remove:
del batch_processors[batch_id]
del batch_processors[batch_id]
logger.info(f"Cleaned up old batch processor: {batch_id}")

72
backend/config.py Normal file
View file

@ -0,0 +1,72 @@
#!/usr/bin/env python3
"""
Configuration settings for DocTags
"""
import os
from pathlib import Path
# Application settings
APP_NAME = "DocTags Intelligence Suite"
APP_VERSION = "1.0.0"
DEBUG = os.environ.get('DEBUG', 'False').lower() == 'true'
# Server settings
HOST = '127.0.0.1'
PORT = 5000
MAX_CONTENT_LENGTH = 100 * 1024 * 1024 # 100MB
# Processing settings
DEFAULT_DPI = 200
PREVIEW_DPI = 150
DEFAULT_GRID_SIZE = 500
MAX_IMAGE_WIDTH = 1200
DEFAULT_PAGE = 1
PROCESSING_TIMEOUT = 300 # 5 minutes
BATCH_WORKERS = 4
# File settings
ALLOWED_EXTENSIONS = {'pdf'}
RESULTS_DIR = 'results'
UPLOAD_DIR = 'uploads'
TEMP_DIR = 'temp_uploads'
# Cleanup settings
CLEANUP_AGE_HOURS = 24
CLEANUP_INTERVAL = 3600 # 1 hour
# Model settings
MODEL_PATH = "ds4sd/SmolDocling-256M-preview-mlx-bf16"
MAX_TOKENS = 4096
# Zone colors for visualization
ZONE_COLORS = {
'section_header_level_1': (255, 87, 34), # Orange
'text': (33, 150, 243), # Blue
'picture': (76, 175, 80), # Green
'table': (156, 39, 176), # Purple
'page_header': (255, 193, 7), # Amber
'page_footer': (121, 85, 72), # Brown
'default': (96, 125, 139) # Blue Grey
}
# Logging configuration
LOGGING_CONFIG = {
'version': 1,
'disable_existing_loggers': False,
'formatters': {
'default': {
'format': '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
},
},
'handlers': {
'console': {
'class': 'logging.StreamHandler',
'formatter': 'default',
},
},
'root': {
'level': 'INFO',
'handlers': ['console'],
},
}

440
backend/page_treatment/analyzer.py
View file

@ -8,39 +8,27 @@
# "requests",
# "argparse",
# "pdf2image",
# "pymupdf", # Optional for better PDF handling
# ]
# ///
import argparse
import os
import tempfile
import re
import base64
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse
import requests
from PIL import Image, UnidentifiedImageError
from pdf2image import convert_from_path, convert_from_bytes
from PIL import Image
from pdf2image import convert_from_bytes
from docling_core.types.doc import ImageRefMode
from docling_core.types.doc.document import DocTagsDocument, DoclingDocument
# Add parent directory to path for imports
import sys
sys.path.append(str(Path(__file__).parent.parent.parent))
def ensure_results_folder():
"""Create the results folder if it doesn't exist."""
# Get the project root directory (where the script is called from)
# Since we're in backend/page_treatment/, we need to go up to the root
script_dir = Path(__file__).parent
project_root = script_dir.parent.parent # Go up two levels from backend/page_treatment/
results_dir = project_root / "results"
if not results_dir.exists():
results_dir.mkdir(parents=True)
print(f"Created results directory: {results_dir}")
print(f"Using results directory: {results_dir.absolute()}")
return results_dir
from backend.utils import ensure_results_folder, load_pdf_page, get_project_root
from backend.config import MODEL_PATH, MAX_TOKENS, DEFAULT_DPI
def parse_arguments():
"""Parse command line arguments."""
@ -53,257 +41,69 @@ def parse_arguments():
help='Prompt for the model')
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "output.html"),
help='Output file path')
parser.add_argument('--show', '-s', action='store_true',
help='Show output in browser')
parser.add_argument('--page', type=int, default=1,
help='Page number to process for PDF files (starts at 1)')
parser.add_argument('--dpi', type=int, default=200,
help='DPI for PDF rendering (higher values produce larger images)')
parser.add_argument('--debug', '-d', action='store_true',
help='Enable debug mode with extra output')
parser.add_argument('--doctags-only', action='store_true',
help='Generate only raw DocTags output without processing')
parser.add_argument('--all-pages', '-a', action='store_true',
help='Process all pages in a PDF without asking')
parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering')
parser.add_argument('--start-page', type=int, default=1,
help='Start processing PDF from this page number')
parser.add_argument('--end-page', type=int, default=None,
help='Stop processing PDF at this page number')
parser.add_argument('--max-pages', type=int, default=None,
help='Maximum number of pages to process')
return parser.parse_args()
def load_image(image_path, page_num=1, dpi=200):
def load_image(image_path, page_num=1, dpi=DEFAULT_DPI):
"""Load image from URL, local image file, or PDF."""
if urlparse(image_path).scheme in ['http', 'https']: # it is a URL
try:
response = requests.get(image_path, stream=True, timeout=10)
response.raise_for_status()
content = response.content
if urlparse(image_path).scheme in ['http', 'https']:
response = requests.get(image_path, stream=True, timeout=10)
response.raise_for_status()
# Check if it's a PDF
if image_path.lower().endswith('.pdf') or response.headers.get('Content-Type') == 'application/pdf':
print(f"Converting PDF from URL (page {page_num})...")
pdf_images = convert_from_bytes(content, dpi=dpi, first_page=page_num, last_page=page_num)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0] # Return the first (and only) page
else:
return Image.open(BytesIO(content))
except requests.exceptions.RequestException as e:
raise Exception(f"Error loading image from URL: {e}")
else: # it is a local file
if image_path.lower().endswith('.pdf') or response.headers.get('Content-Type') == 'application/pdf':
print(f"Converting PDF from URL (page {page_num})...")
pdf_images = convert_from_bytes(response.content, dpi=dpi, first_page=page_num, last_page=page_num)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0]
else:
return Image.open(response.raw)
else:
image_path = Path(image_path)
if not image_path.exists():
raise FileNotFoundError(f"File not found: {image_path}")
# Check if it's a PDF
if image_path.suffix.lower() == '.pdf':
print(f"Converting PDF to image (page {page_num}, DPI: {dpi})...")
try:
pdf_images = convert_from_path(
image_path,
dpi=dpi,
first_page=page_num,
last_page=page_num
)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0] # Return the requested page
except Exception as e:
raise Exception(f"Error converting PDF to image: {e}")
return load_pdf_page(str(image_path), page_num, dpi)
else:
try:
return Image.open(image_path)
except UnidentifiedImageError:
raise Exception(f"Cannot identify image file: {image_path}. Make sure it's a valid image format or PDF.")
return Image.open(image_path)
def cleanup_doctags(doctags_text):
"""Clean up the DocTags structure."""
print("Cleaning up DocTags structure...")
# Simplified approach to extract valuable information
# Extract headers
headers = re.findall(r'<section_header_level_1>.*?>(.*?)</section_header_level_1>', doctags_text)
# Extract text paragraphs
paragraphs = re.findall(r'<text>.*?>(.*?)</text>', doctags_text)
# Extract list items
list_items = re.findall(r'<list_item>.*?>(.*?)</list_item>', doctags_text)
# Extract footer
footer = re.search(r'<page_footer>.*?>(.*?)</page_footer>', doctags_text)
footer_text = footer.group(1) if footer else ""
# Create a clean doctags structure
clean_doctags = "<doctag>\n"
# Add headers
for header in headers:
clean_doctags += f"<section_header_level_1>{header}</section_header_level_1>\n"
# Add text
for paragraph in paragraphs:
clean_doctags += f"<text>{paragraph}</text>\n"
# Add list if any items found
if list_items:
clean_doctags += "<unordered_list>\n"
for item in list_items:
clean_doctags += f"<list_item>{item}</list_item>\n"
clean_doctags += "</unordered_list>\n"
# Add footer if present
if footer_text:
clean_doctags += f"<page_footer>{footer_text}</page_footer>\n"
clean_doctags += "</doctag>"
return clean_doctags
def extract_all_tags(doctags_text):
"""Extract all unique DocTags from the text."""
print("Extracting all DocTags...")
# Use regex to find all tags
tag_pattern = r'</?(\w+)(?:\s[^>]*)?>'
all_tags = re.findall(tag_pattern, doctags_text)
# Remove duplicates and sort
unique_tags = sorted(set(all_tags))
# Create a list of tags with their frequencies
tag_counts = {}
for tag in all_tags:
tag_counts[tag] = tag_counts.get(tag, 0) + 1
# Format the output
tags_output = "# DocTags Found\n\n"
tags_output += "| Tag | Count |\n"
tags_output += "|-----|-------|\n"
for tag in unique_tags:
tags_output += f"| {tag} | {tag_counts[tag]} |\n"
# Add examples section
tags_output += "\n\n# DocTags Examples\n\n"
# Find example usages for each tag
for tag in unique_tags:
# Find an opening tag with content
open_pattern = f'<{tag}(?:\\s[^>]*)?>.*?</{tag}>'
examples = re.findall(open_pattern, doctags_text, re.DOTALL)
if examples:
# Limit to first example
example = examples[0]
# Truncate if too long
if len(example) > 200:
example = example[:197] + "..."
tags_output += f"## {tag}\n\n```xml\n{example}\n```\n\n"
return tags_output
def debug_doctags(doctags_text, debug_mode=False):
"""Debug function to analyze doctag structure."""
if not debug_mode:
return doctags_text
print("\nAnalyzing DocTags structure:")
print(f"Total length: {len(doctags_text)} characters")
# Check for valid opening and closing tags
opening_tags = []
i = 0
while i < len(doctags_text):
open_tag_start = doctags_text.find('<', i)
if open_tag_start == -1:
break
open_tag_end = doctags_text.find('>', open_tag_start)
if open_tag_end == -1:
print(f"WARNING: Unclosed tag starting at position {open_tag_start}")
break
tag_content = doctags_text[open_tag_start+1:open_tag_end]
if tag_content.startswith('/'):
# This is a closing tag
if not opening_tags:
print(f"WARNING: Closing tag {tag_content} without matching opening tag")
else:
last_open = opening_tags.pop()
if last_open != tag_content[1:]:
print(f"WARNING: Mismatched tags: opening <{last_open}> vs closing <{tag_content}>")
elif not tag_content.startswith('!') and not ' ' in tag_content:
# This is an opening tag (not a comment or self-closing)
opening_tags.append(tag_content)
i = open_tag_end + 1
if opening_tags:
print(f"WARNING: Unclosed tags: {', '.join(opening_tags)}")
# Count all tags
tag_counts = {}
tag_pattern = r'<(\w+)(?:\s|>)'
import re
for tag in re.findall(tag_pattern, doctags_text):
tag_counts[tag] = tag_counts.get(tag, 0) + 1
print("Tag counts:")
for tag, count in tag_counts.items():
print(f" <{tag}>: {count}")
# Special check for doctag
if doctags_text.count('<doctag>') != 1:
print(f"WARNING: Expected 1 <doctag> tag, found {doctags_text.count('<doctag>')}")
if doctags_text.count('</doctag>') != 1:
print(f"WARNING: Expected 1 </doctag> tag, found {doctags_text.count('</doctag>')}")
return doctags_text
def process_page(args, model, processor, config, image_path, pil_image, page_num=1):
def process_page(model, processor, config, args, pil_image, page_num=1):
"""Process a single page from a PDF or image file."""
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import stream_generate
# Ensure results folder exists
results_dir = ensure_results_folder()
# Prepare input
prompt = args.prompt
output_base = Path(args.output)
output_path = output_base
# If processing PDF and output is a path without explicit numbering, add page numbers
if Path(image_path).suffix.lower() == '.pdf' and page_num > 1:
# Get base filename without extension
base_name = output_base.stem
output_path = results_dir / f"{base_name}_page{page_num}{output_base.suffix}"
# Handle multi-page output naming
if Path(args.image).suffix.lower() == '.pdf' and page_num > 1:
output_path = results_dir / f"{output_base.stem}_page{page_num}{output_base.suffix}"
else:
output_path = output_base
print(f"Processing page {page_num}, output will be saved to {output_path}")
# Create a temporary file for the image
# Save image temporarily
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img_file:
temp_img_path = temp_img_file.name
pil_image.save(temp_img_path, format='PNG')
print(f"Saved temporary image to: {temp_img_path}")
try:
# Apply chat template
formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1)
# Apply chat template and generate
formatted_prompt = apply_chat_template(processor, config, args.prompt, num_images=1)
# Generate output
print(f"Generating DocTags for page {page_num}: \n\n")
output = ""
for token in stream_generate(
model, processor, formatted_prompt, [temp_img_path], max_tokens=4096, verbose=False
model, processor, formatted_prompt, [temp_img_path], max_tokens=MAX_TOKENS, verbose=False
):
output += token.text
print(token.text, end="")
@ -311,184 +111,52 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
break
print("\n\n")
# Debug and clean the doctags content
output = debug_doctags(output, args.debug)
# Extract all tags if in DocTags-only mode
if args.doctags_only:
tags_analysis = extract_all_tags(output)
# Clean the output for document creation
cleaned_output = cleanup_doctags(output)
finally:
# Clean up the temporary file
# Clean up temporary file
if os.path.exists(temp_img_path):
os.unlink(temp_img_path)
print(f"Removed temporary image file")
# Save the raw DocTags to a txt file in results folder
# Save DocTags output
doctags_path = results_dir / f"{output_path.stem}.doctags.txt"
with open(doctags_path, 'w', encoding='utf-8') as f:
f.write(output)
print(f"Raw DocTags saved to: {doctags_path}")
# Save the tag analysis if in DocTags-only mode
if args.doctags_only:
tags_path = results_dir / f"{output_path.stem}.tags.md"
with open(tags_path, 'w', encoding='utf-8') as f:
f.write(tags_analysis)
print(f"DocTags analysis saved to: {tags_path}")
# Save a copy of the processed image for reference if in debug mode
if args.debug:
img_debug_path = results_dir / f"{output_path.stem}.debug.png"
pil_image.save(img_debug_path)
print(f"Saved debug image to: {img_debug_path}")
return output_path
def main():
# Ensure results folder exists
ensure_results_folder()
# Parse arguments
args = parse_arguments()
# Settings
DEBUG_MODE = args.debug
DOCTAGS_ONLY = args.doctags_only
image_path = args.image
# Load the model
print("Loading model...")
try:
from mlx_vlm import load, generate
from mlx_vlm import load
from mlx_vlm.utils import load_config
model_path = "ds4sd/SmolDocling-256M-preview-mlx-bf16"
model, processor = load(model_path)
config = load_config(model_path)
model, processor = load(MODEL_PATH)
config = load_config(MODEL_PATH)
except Exception as e:
print(f"Error loading model: {e}")
import traceback
traceback.print_exc()
return
# Check if the input is a PDF
pdf_path = Path(image_path)
is_pdf = False
pdf_page_count = 0
# Process the image/PDF
try:
# Handle single page or range
start_page = args.start_page
end_page = args.end_page or args.page
if pdf_path.suffix.lower() == '.pdf':
is_pdf = True
# Try to get page count
try:
try:
import fitz # PyMuPDF
pdf_document = fitz.open(pdf_path)
pdf_page_count = len(pdf_document)
print(f"PDF detected with {pdf_page_count} pages")
pdf_document.close()
except ImportError:
print("PyMuPDF not installed. Using pdf2image to estimate page count...")
from pdf2image import pdfinfo_from_path
pdf_info = pdfinfo_from_path(pdf_path)
pdf_page_count = pdf_info["Pages"]
print(f"PDF detected with {pdf_page_count} pages")
except Exception as e:
print(f"Could not determine PDF page count: {e}")
print("Will process the specified page only.")
pdf_page_count = args.page
for page_num in range(start_page, end_page + 1):
print(f"\nProcessing page {page_num}...")
# Determine which pages to process
process_all_pages = args.all_pages
start_page = args.start_page
end_page = args.end_page if args.end_page else pdf_page_count
max_pages = args.max_pages
pil_image = load_image(args.image, page_num=page_num, dpi=args.dpi)
print(f"Page {page_num} loaded: {pil_image.size}")
# Validate page ranges
if start_page < 1:
start_page = 1
if end_page and end_page > pdf_page_count:
end_page = pdf_page_count
# Ask user if they want to process all pages or just one (if not specified by arguments)
if is_pdf and pdf_page_count > 1 and not process_all_pages and args.start_page == 1 and not args.end_page:
if args.page > 1:
print(f"You specified to process page {args.page}.")
process_all_pages = False
start_page = args.page
end_page = args.page
else:
user_input = input("Do you want to process all pages? [y/N]: ")
process_all_pages = user_input.lower() in ['y', 'yes']
if not process_all_pages:
# Ask for specific page or range
page_input = input(f"Enter page number(s) to process (e.g., 3 or 1-5) [1]: ")
if page_input.strip():
if '-' in page_input:
try:
start_str, end_str = page_input.split('-')
start_page = int(start_str.strip())
end_page = int(end_str.strip())
except ValueError:
print("Invalid range format. Using default page 1.")
start_page = end_page = 1
else:
try:
start_page = end_page = int(page_input.strip())
except ValueError:
print("Invalid page number. Using default page 1.")
start_page = end_page = 1
# Apply max_pages limit if specified
if max_pages and end_page - start_page + 1 > max_pages:
end_page = start_page + max_pages - 1
# Process pages
if is_pdf and (process_all_pages or start_page != end_page):
page_range = range(start_page, end_page + 1)
print(f"Processing pages {start_page} to {end_page} ({len(page_range)} pages)...")
processed_pages = []
for page_num in page_range:
print(f"\n{'='*50}\nProcessing PDF page {page_num}/{end_page}\n{'='*50}\n")
# Update the page argument
args.page = page_num
# Load the specific page
try:
pil_image = load_image(image_path, page_num=page_num, dpi=args.dpi)
print(f"Page {page_num} loaded: {pil_image.size}")
# Process the page
output_path = process_page(args, model, processor, config, image_path, pil_image, page_num)
processed_pages.append(output_path)
except Exception as e:
print(f"Error processing page {page_num}: {e}")
import traceback
traceback.print_exc()
continue
print(f"\nProcessed {len(processed_pages)} pages from PDF.")
print(f"Output files: {', '.join([str(p) for p in processed_pages])}")
else:
# Process just one page (either it's not a PDF or user only wants one page)
try:
# Load image resource
print(f"Loading {'PDF page' if is_pdf else 'image'} from: {image_path}")
pil_image = load_image(image_path, page_num=args.page, dpi=args.dpi)
print(f"Image loaded: {pil_image.size}")
# Process the single page
process_page(args, model, processor, config, image_path, pil_image)
except Exception as e:
print(f"Error processing {image_path}: {e}")
import traceback
traceback.print_exc()
process_page(model, processor, config, args, pil_image, page_num)
except Exception as e:
print(f"Error processing: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()

259
backend/page_treatment/picture_extractor.py
View file

@ -1,41 +1,26 @@
#!/usr/bin/env python3
"""
DocTags Picture Extractor - Extract <picture> elements from DocTags and save as separate image files.
Usage:
python picture_extractor.py --doctags output.doctags.txt --pdf document.pdf --page 1 --output pictures
DocTags Picture Extractor - Extract <picture> elements from DocTags.
"""
import argparse
import os
import re
import sys
from io import BytesIO
from pathlib import Path
import pdf2image
from PIL import Image
# Add parent directory to path for imports
import sys
sys.path.append(str(Path(__file__).parent.parent.parent))
from backend.utils import (ensure_results_folder, load_pdf_page,
normalize_coordinates, auto_adjust_coordinates,
validate_coordinates)
from backend.config import DEFAULT_DPI, MAX_IMAGE_WIDTH, DEFAULT_GRID_SIZE
# Regular expression to extract picture location data
PICTURE_PATTERN = r'<picture>.*?<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>(.*?)</picture>'
def ensure_results_folder(custom_path=None):
"""Create the results folder if it doesn't exist."""
if custom_path:
results_dir = Path(custom_path)
else:
# Get the project root directory (where the script is called from)
# Since we're in backend/page_treatment/, we need to go up to the root
script_dir = Path(__file__).parent
project_root = script_dir.parent.parent # Go up two levels from backend/page_treatment/
results_dir = project_root / "results"
if not results_dir.exists():
results_dir.mkdir(parents=True)
print(f"Created directory: {results_dir}")
print(f"Using results directory: {results_dir.absolute()}")
return results_dir
def parse_arguments():
"""Parse command line arguments."""
results_dir = ensure_results_folder()
@ -46,46 +31,19 @@ def parse_arguments():
parser.add_argument('--pdf', '-p', type=str, required=True,
help='Path to original PDF file')
parser.add_argument('--page', type=int, default=1,
help='Page number in PDF (starts at 1, default: 1)')
help='Page number in PDF (starts at 1)')
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "pictures"),
help='Output directory for extracted pictures')
parser.add_argument('--dpi', type=int, default=300,
help='DPI for PDF rendering (higher values produce larger images)')
parser.add_argument('--max-width', type=int, default=1200,
parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering')
parser.add_argument('--max-width', type=int, default=MAX_IMAGE_WIDTH,
help='Maximum width of output images in pixels')
parser.add_argument('--adjust', action='store_true',
help='Try to automatically adjust scaling')
parser.add_argument('--scale', type=float, default=1.0,
help='Scaling factor for coordinates (default: 1.0)')
parser.add_argument('--scale-x', type=float, default=None,
help='X-axis scaling factor (overrides --scale)')
parser.add_argument('--scale-y', type=float, default=None,
help='Y-axis scaling factor (overrides --scale)')
parser.add_argument('--margin', type=int, default=0,
help='Add margin around extracted pictures in pixels')
parser.add_argument('--show', '-s', action='store_true',
help='Open a file browser to the output directory when done')
return parser.parse_args()
def load_image_from_pdf(pdf_path, page_num=1, dpi=300):
"""Load a specific page from PDF as an image."""
if not os.path.exists(pdf_path):
raise FileNotFoundError(f"PDF file not found: {pdf_path}")
print(f"Converting PDF page {page_num} to image (DPI: {dpi})...")
try:
pdf_images = pdf2image.convert_from_path(
pdf_path,
dpi=dpi,
first_page=page_num,
last_page=page_num
)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0] # Return the requested page
except Exception as e:
raise Exception(f"Error converting PDF to image: {e}")
def extract_pictures_from_doctags(doctags_path):
"""Parse DocTags file and extract picture elements with their coordinates."""
if not os.path.exists(doctags_path):
@ -94,116 +52,40 @@ def extract_pictures_from_doctags(doctags_path):
with open(doctags_path, 'r', encoding='utf-8') as f:
doctags_content = f.read()
# Find all picture elements with location information
pictures = []
picture_matches = re.finditer(PICTURE_PATTERN, doctags_content, re.DOTALL)
for i, match in enumerate(picture_matches):
x1, y1, x2, y2, caption = match.groups()
# Extract caption if available (remove location tags)
# Clean caption
clean_caption = re.sub(r'<loc_\d+>', '', caption).strip()
pictures.append({
'id': i + 1,
'x1': int(x1),
'y1': int(y1),
'x2': int(x2),
'y2': int(y2),
'x1': int(x1), 'y1': int(y1),
'x2': int(x2), 'y2': int(y2),
'caption': clean_caption
})
return pictures
def normalize_coordinates(pictures, image_width, image_height, grid_size=500):
"""
Normalize coordinates from the DocTags grid (0-500) to actual image dimensions.
"""
normalized_pictures = []
for picture in pictures:
# Clone the picture
new_picture = picture.copy()
# Convert from grid coordinates to actual page dimensions
new_picture['x1'] = int(picture['x1'] * image_width / grid_size)
new_picture['y1'] = int(picture['y1'] * image_height / grid_size)
new_picture['x2'] = int(picture['x2'] * image_width / grid_size)
new_picture['y2'] = int(picture['y2'] * image_height / grid_size)
normalized_pictures.append(new_picture)
return normalized_pictures
def auto_adjust_coordinates(pictures, image_width, image_height):
"""
Automatically adjust coordinates based on image dimensions.
"""
if not pictures:
return pictures
# Find the maximum coordinates
max_x = max([pic['x2'] for pic in pictures])
max_y = max([pic['y2'] for pic in pictures])
# If coordinates seem to be in a normalized grid (0-500 range)
if max_x <= 500 and max_y <= 500:
print(f"Detected normalized coordinates (0-500 grid)")
return normalize_coordinates(pictures, image_width, image_height)
# Calculate appropriate scaling factors with better heuristics
if max_x > 0:
x_scale = min(image_width / max_x, 1.0) if max_x > image_width else max(image_width / max_x, 0.5)
print(f"Auto-adjusted X scale to {x_scale:.3f} (image width: {image_width}, max picture x: {max_x})")
else:
x_scale = 1.0
if max_y > 0:
y_scale = min(image_height / max_y, 1.0) if max_y > image_height else max(image_height / max_y, 0.5)
print(f"Auto-adjusted Y scale to {y_scale:.3f} (image height: {image_height}, max picture y: {max_y})")
else:
y_scale = 1.0
# Apply more aggressive adjustment if image and pictures are very different in scale
if max_x > image_width * 5 or max_x < image_width / 5:
x_scale = image_width / max_x
print(f"Major X scale adjustment to {x_scale:.3f}")
if max_y > image_height * 5 or max_y < image_height / 5:
y_scale = image_height / max_y
print(f"Major Y scale adjustment to {y_scale:.3f}")
# Apply the scaling to all pictures
adjusted_pictures = []
for pic in pictures:
adjusted_pic = pic.copy()
adjusted_pic['x1'] = int(pic['x1'] * x_scale)
adjusted_pic['y1'] = int(pic['y1'] * y_scale)
adjusted_pic['x2'] = int(pic['x2'] * x_scale)
adjusted_pic['y2'] = int(pic['y2'] * y_scale)
adjusted_pictures.append(adjusted_pic)
print(f"Applied auto-scaling: X={x_scale}, Y={y_scale}")
return adjusted_pictures
def extract_and_save_pictures(image, pictures, output_dir, max_width=1200, margin=0):
def extract_and_save_pictures(image, pictures, output_dir, max_width, margin):
"""Extract picture regions from the image and save them as separate files."""
# Ensure output directory exists
output_path = ensure_results_folder(output_dir)
saved_files = []
# Process each picture
for picture in pictures:
try:
# Add margin to coordinates if specified
# Add margin to coordinates
x1 = max(0, picture['x1'] - margin)
y1 = max(0, picture['y1'] - margin)
x2 = min(image.width, picture['x2'] + margin)
y2 = min(image.height, picture['y2'] + margin)
# Check if coordinates are valid
if x1 >= x2 or y1 >= y2 or x1 < 0 or y1 < 0 or x2 > image.width or y2 > image.height:
print(f"Warning: Invalid coordinates for picture {picture['id']}: ({x1},{y1})-({x2},{y2})")
# Validate coordinates
if not validate_coordinates(x1, y1, x2, y2, image.width, image.height):
print(f"Warning: Invalid coordinates for picture {picture['id']}")
continue
# Crop the image
@ -216,10 +98,8 @@ def extract_and_save_pictures(image, pictures, output_dir, max_width=1200, margi
cropped_img = cropped_img.resize((max_width, new_height), Image.LANCZOS)
# Generate filename
caption = picture['caption']
if caption:
# Create a filename-safe version of the caption (first 30 chars)
safe_caption = re.sub(r'[^\w\s-]', '', caption)[:30].strip().replace(' ', '_').lower()
if picture['caption']:
safe_caption = re.sub(r'[^\w\s-]', '', picture['caption'])[:30].strip().replace(' ', '_').lower()
filename = f"picture_{picture['id']}_{safe_caption}.png"
else:
filename = f"picture_{picture['id']}.png"
@ -228,11 +108,11 @@ def extract_and_save_pictures(image, pictures, output_dir, max_width=1200, margi
output_file = output_path / filename
cropped_img.save(output_file, format="PNG")
# Create a text file with the caption if available
if caption:
# Save caption if available
if picture['caption']:
caption_file = output_path / f"{output_file.stem}.txt"
with open(caption_file, 'w', encoding='utf-8') as f:
f.write(caption)
f.write(picture['caption'])
print(f"Saved picture {picture['id']} to {output_file}")
saved_files.append(output_file)
@ -275,20 +155,10 @@ def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir):
.picture-info {{
padding: 15px;
}}
.picture-caption {{
margin-top: 10px;
color: #555;
}}
.picture-coords {{
margin-top: 5px;
font-size: 0.8em;
color: #777;
}}
.no-pictures {{
background-color: white;
padding: 20px;
border-radius: 5px;
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
text-align: center;
color: #777;
}}
@ -300,42 +170,26 @@ def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir):
"""
if pictures:
html += """ <div class="gallery">
"""
html += ' <div class="gallery">\n'
for picture, file_path in zip(pictures, saved_files):
# Get relative path for the image
rel_path = file_path.name
html += f""" <div class="picture-card">
<img src="{rel_path}" alt="Picture {picture['id']}">
<div class="picture-info">
<h3>Picture {picture['id']}</h3>
"""
if picture['caption']:
html += f""" <div class="picture-caption">{picture['caption']}</div>
"""
html += f""" <div class="picture-coords">Coordinates: ({picture['x1']},{picture['y1']})-({picture['x2']},{picture['y2']})</div>
{f'<div class="picture-caption">{picture["caption"]}</div>' if picture['caption'] else ''}
<div class="picture-coords">Coordinates: ({picture['x1']},{picture['y1']})-({picture['x2']},{picture['y2']})</div>
</div>
</div>
"""
html += """ </div>
"""
html += ' </div>\n'
else:
html += """ <div class="no-pictures">
<h2>No pictures found on this page</h2>
<p>The DocTags file doesn't contain any picture elements for this page.</p>
</div>
"""
html += ' <div class="no-pictures">\n <h2>No pictures found on this page</h2>\n </div>\n'
html += """</body>
</html>
"""
html += '</body>\n</html>\n'
# Save the HTML file
with open(index_file, 'w', encoding='utf-8') as f:
f.write(html)
@ -343,10 +197,7 @@ def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir):
return index_file
def main():
# Parse arguments
args = parse_arguments()
# Create output directory
output_dir = ensure_results_folder(args.output)
try:
@ -361,48 +212,30 @@ def main():
print(f"Found {len(pictures)} picture elements.")
# Load the image from PDF
page_image = load_image_from_pdf(args.pdf, args.page, args.dpi)
page_image = load_pdf_page(args.pdf, args.page, args.dpi)
print(f"Loaded page {args.page} image: {page_image.size[0]}x{page_image.size[1]}")
# Process coordinates
# Adjust coordinates if needed
if args.adjust:
pictures = auto_adjust_coordinates(pictures, page_image.width, page_image.height)
elif args.scale != 1.0 or args.scale_x is not None or args.scale_y is not None:
# Apply manual scaling
scale_x = args.scale_x if args.scale_x is not None else args.scale
scale_y = args.scale_y if args.scale_y is not None else args.scale
# Check if coordinates need normalization
max_x = max([p['x2'] for p in pictures])
max_y = max([p['y2'] for p in pictures])
print(f"Applying manual scaling: X={scale_x}, Y={scale_y}")
for picture in pictures:
picture['x1'] = int(picture['x1'] * scale_x)
picture['y1'] = int(picture['y1'] * scale_y)
picture['x2'] = int(picture['x2'] * scale_x)
picture['y2'] = int(picture['y2'] * scale_y)
if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
print(f"Detected normalized coordinates (0-{DEFAULT_GRID_SIZE} grid)")
pictures = normalize_coordinates(pictures, page_image.width, page_image.height)
else:
pictures = auto_adjust_coordinates(pictures, page_image.width, page_image.height)
# Extract and save pictures
saved_files = extract_and_save_pictures(
page_image,
pictures,
output_dir,
args.max_width,
args.margin
page_image, pictures, output_dir,
args.max_width, args.margin
)
# Create HTML index
pdf_name = Path(args.pdf).stem
index_file = create_html_index(pictures, saved_files, pdf_name, args.page, output_dir)
# Open the output directory or index file if requested
if args.show and saved_files:
import webbrowser
if sys.platform == 'darwin': # macOS
import subprocess
subprocess.run(['open', str(output_dir)])
elif sys.platform == 'win32': # Windows
import os
os.startfile(str(output_dir))
else: # Linux
webbrowser.open(f"file:///{os.path.abspath(index_file)}")
create_html_index(pictures, saved_files, pdf_name, args.page, output_dir)
except Exception as e:
print(f"Error: {e}")

350
backend/page_treatment/visualizer.py
View file

@ -1,114 +1,44 @@
#!/usr/bin/env python3
"""
DocTags Zone Visualizer - Simple script to visualize zones identified in DocTags format.
PNG-only version: Creates debug images with rectangles around zones.
Usage:
python visualizer.py --doctags output.doctags.txt --pdf document.pdf --page 8
DocTags Zone Visualizer - Visualize zones identified in DocTags format.
"""
import argparse
import os
import re
import sys
from pathlib import Path
from PIL import Image, ImageDraw
import pdf2image
# Add parent directory to path for imports
import sys
sys.path.append(str(Path(__file__).parent.parent.parent))
from backend.utils import (ensure_results_folder, load_pdf_page, count_pdf_pages,
normalize_coordinates, auto_adjust_coordinates)
from backend.config import ZONE_COLORS, DEFAULT_DPI, DEFAULT_GRID_SIZE
# Regular expression to extract location data
LOC_PATTERN = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>'
def ensure_results_folder():
"""Create the results folder if it doesn't exist."""
# Get the project root directory (where the script is called from)
# Since we're in backend/page_treatment/, we need to go up to the root
script_dir = Path(__file__).parent
project_root = script_dir.parent.parent # Go up two levels from backend/page_treatment/
results_dir = project_root / "results"
if not results_dir.exists():
results_dir.mkdir(parents=True)
print(f"Created results directory: {results_dir}")
print(f"Using results directory: {results_dir.absolute()}")
return results_dir
def parse_arguments():
"""Parse command line arguments."""
results_dir = ensure_results_folder()
parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format as PNG images')
parser = argparse.ArgumentParser(description='Visualize zones identified in DocTags format')
parser.add_argument('--doctags', '-d', type=str, required=True,
help='Path to DocTags file')
parser.add_argument('--pdf', '-p', type=str, required=True,
help='Path to original PDF file')
parser.add_argument('--page', type=int, default=8,
help='Page number in PDF (starts at 1, default: 8)')
parser.add_argument('--page', type=int, default=1,
help='Page number in PDF (starts at 1)')
parser.add_argument('--output', '-o', type=str, default=None,
help='Output PNG file path (default: results/visualization_page_X.png)')
parser.add_argument('--dpi', type=int, default=200,
help='Output PNG file path')
parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering')
parser.add_argument('--page-count', action='store_true',
help='Just count pages in the PDF and exit')
parser.add_argument('--scale', type=float, default=1.0,
help='Scaling factor for zone coordinates (default: 1.0)')
parser.add_argument('--scale-x', type=float, default=None,
help='X-axis scaling factor (overrides --scale)')
parser.add_argument('--scale-y', type=float, default=None,
help='Y-axis scaling factor (overrides --scale)')
parser.add_argument('--adjust', action='store_true',
help='Try to automatically adjust scaling')
return parser.parse_args()
def count_pdf_pages(pdf_path):
"""Count the number of pages in a PDF file."""
if not os.path.exists(pdf_path):
print(f"Error: PDF file not found: {pdf_path}")
return 0
try:
from pdf2image.pdf2image import pdfinfo_from_path
info = pdfinfo_from_path(pdf_path, userpw=None, poppler_path=None)
return info["Pages"]
except Exception as e:
print(f"Warning: pdfinfo failed: {e}")
# Fallback method if pdfinfo fails
try:
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=1, last_page=1)
# Try to load the last page - increment until we get an error
page_count = 1
while True:
try:
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=page_count+1, last_page=page_count+1)
if not images:
break
page_count += 1
except:
break
return page_count
except Exception as e2:
print(f"Error counting PDF pages: {e2}")
return 0
def load_image_from_pdf(pdf_path, page_num=1, dpi=200):
"""Load a specific page from PDF as an image."""
if not os.path.exists(pdf_path):
raise FileNotFoundError(f"PDF file not found: {pdf_path}")
print(f"Converting PDF page {page_num} to image (DPI: {dpi})...")
try:
pdf_images = pdf2image.convert_from_path(
pdf_path,
dpi=dpi,
first_page=page_num,
last_page=page_num
)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0] # Return the requested page
except Exception as e:
raise Exception(f"Error converting PDF to image: {e}")
def parse_doctags(doctags_path):
"""Parse DocTags file and extract zones with their coordinates."""
if not os.path.exists(doctags_path):
@ -118,86 +48,59 @@ def parse_doctags(doctags_path):
doctags_content = f.read()
# Extract content between <doctag> tags
doctag_pattern = r'<doctag>(.*?)</doctag>'
doctag_match = re.search(doctag_pattern, doctags_content, re.DOTALL)
doctag_match = re.search(r'<doctag>(.*?)</doctag>', doctags_content, re.DOTALL)
if not doctag_match:
raise ValueError("No <doctag> tags found in the file")
doctag_content = doctag_match.group(1)
# Find all tags with location information
zones = []
# Find all tag starts
# Find all tags with location information
tag_starts = re.finditer(r'<(\w+)>', doctag_content)
for tag_match in tag_starts:
tag_name = tag_match.group(1)
# Skip location tags themselves
if tag_name.startswith('loc_'):
continue
# Find the end of the tag
tag_start_pos = tag_match.start()
tag_end_pattern = f'</({tag_name})>'
tag_end_match = re.search(tag_end_pattern, doctag_content[tag_start_pos:])
if not tag_end_match:
continue # Skip if no closing tag
continue
# Extract the tag content
tag_content = doctag_content[tag_start_pos:tag_start_pos + tag_end_match.end()]
# Look for location pattern
loc_match = re.search(LOC_PATTERN, tag_content)
if loc_match:
# Extract coordinates
x1, y1, x2, y2 = map(int, loc_match.groups())
# Extract text content if available
text_content = ""
# Look for content between the location info and the closing tag
# Extract text content
content_pattern = f'{LOC_PATTERN}(.*?)</{tag_name}>'
content_match = re.search(content_pattern, tag_content, re.DOTALL)
if content_match:
text_content = content_match.group(5).strip()
text_content = content_match.group(5).strip() if content_match else ""
zones.append({
'type': tag_name,
'x1': x1,
'y1': y1,
'x2': x2,
'y2': y2,
'x1': x1, 'y1': y1,
'x2': x2, 'y2': y2,
'content': text_content
})
return zones
def create_debug_image(image, zones, page_num, output_path):
"""Create a debug image with rectangles around zones."""
# Create a copy of the input image
def create_visualization(image, zones, page_num, output_path):
"""Create a visualization image with rectangles around zones."""
debug_img = image.copy()
draw = ImageDraw.Draw(debug_img)
# Define colors for different zone types
zone_colors = {
'section_header_level_1': (255, 87, 34), # Orange
'text': (33, 150, 243), # Blue
'picture': (76, 175, 80), # Green
'table': (156, 39, 176), # Purple
'page_header': (255, 193, 7), # Amber
'page_footer': (121, 85, 72), # Brown
'default': (96, 125, 139) # Blue Grey
}
# Draw rectangles for each zone
for zone in zones:
zone_type = zone['type']
color = zone_colors.get(zone_type, zone_colors['default'])
color = ZONE_COLORS.get(zone_type, ZONE_COLORS['default'])
# Draw rectangle
draw.rectangle(
[(zone['x1'], zone['y1']), (zone['x2'], zone['y2'])],
outline=color,
@ -206,7 +109,7 @@ def create_debug_image(image, zones, page_num, output_path):
# Add zone type label
label_width = len(zone_type) * 7 + 6
label_x = min(zone['x1'], image.width - label_width) # Keep label on image
label_x = min(zone['x1'], image.width - label_width)
draw.rectangle(
[(label_x, zone['y1']), (label_x + label_width, zone['y1'] + 20)],
@ -219,7 +122,7 @@ def create_debug_image(image, zones, page_num, output_path):
fill=color
)
# Draw page number on the debug image
# Draw page number
draw.rectangle(
[(10, 10), (100, 40)],
fill=(0, 0, 0, 180),
@ -231,167 +134,49 @@ def create_debug_image(image, zones, page_num, output_path):
fill=(255, 255, 255)
)
# Save the debug image
# Save the image
debug_img.save(output_path)
print(f"Debug image saved to: {output_path}")
print(f"Absolute path: {output_path.absolute()}")
print(f"Visualization saved to: {output_path}")
return debug_img
def normalize_coordinates(zones, image_width, image_height, grid_size=500):
"""
Normalize coordinates from the DocTags grid (0-500) to actual image dimensions.
Args:
zones: List of zone dictionaries with x1, y1, x2, y2 coordinates
image_width: Width of the PDF page image in pixels
image_height: Height of the PDF page image in pixels
grid_size: The grid size used in DocTags (default 500)
Returns:
The same zones list with updated coordinates
"""
# Create a copy of the zones to avoid modifying the original
normalized_zones = []
for zone in zones:
# Clone the zone
new_zone = zone.copy()
# Convert from grid coordinates to actual page dimensions
new_zone['x1'] = int(zone['x1'] * image_width / grid_size)
new_zone['y1'] = int(zone['y1'] * image_height / grid_size)
new_zone['x2'] = int(zone['x2'] * image_width / grid_size)
new_zone['y2'] = int(zone['y2'] * image_height / grid_size)
normalized_zones.append(new_zone)
return normalized_zones
def process_page(pdf_path, page_num, doctags_path, output_path, dpi=200, scale=1.0, scale_x=None, scale_y=None, adjust=True):
def process_page(pdf_path, page_num, doctags_path, output_path, dpi, adjust):
"""Process a single page of the PDF with visualization."""
# Ensure results folder exists
results_dir = ensure_results_folder()
# Generate output path if not provided
if output_path is None:
output_name = f"visualization_page_{page_num}.png"
output_path = results_dir / output_name
# Make sure output_path is a Path object
output_path = Path(output_path)
output_path = results_dir / f"visualization_page_{page_num}.png"
else:
output_path = Path(output_path)
# Load the page image
try:
image = load_image_from_pdf(pdf_path, page_num, dpi)
print(f"Page {page_num} loaded: {image.size}")
except Exception as e:
print(f"Error loading page {page_num}: {e}")
return False
image = load_pdf_page(pdf_path, page_num, dpi)
print(f"Page {page_num} loaded: {image.size}")
# Parse DocTags
try:
zones = parse_doctags(doctags_path)
print(f"Found {len(zones)} zones in DocTags")
zones = parse_doctags(doctags_path)
print(f"Found {len(zones)} zones in DocTags")
# Debug output to understand scaling issues
# After parsing the zones from DocTags
if zones:
img_width, img_height = image.size
print(f"Image dimensions: {img_width}x{img_height}")
if zones:
# Check if we need to adjust coordinates
max_x = max([zone['x2'] for zone in zones])
max_y = max([zone['y2'] for zone in zones])
# Check if we need to normalize grid coordinates
max_x = max([zone['x2'] for zone in zones])
max_y = max([zone['y2'] for zone in zones])
# Auto-adjust if needed
if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
print(f"Detected normalized coordinates (0-{DEFAULT_GRID_SIZE} grid)")
zones = normalize_coordinates(zones, image.width, image.height)
elif adjust:
zones = auto_adjust_coordinates(zones, image.width, image.height)
# If coordinates seem to be in a normalized grid (0-500 range)
if max_x <= 500 and max_y <= 500:
print(f"Detected normalized coordinates (0-500 grid)")
zones = normalize_coordinates(zones, img_width, img_height)
print(f"Applied automatic grid normalization")
# If auto-adjust is enabled and coordinates are not in normalized grid
elif adjust:
width, height = image.size
# Calculate appropriate scaling factors with better heuristics
# Use smaller scaling to avoid cutting off content
if max_x > 0:
x_scale = min(width / max_x, 1.0) if max_x > width else max(width / max_x, 0.5)
print(f"Auto-adjusted X scale to {x_scale:.3f} (image width: {width}, max zone x: {max_x})")
else:
x_scale = 1.0
if max_y > 0:
y_scale = min(height / max_y, 1.0) if max_y > height else max(height / max_y, 0.5)
print(f"Auto-adjusted Y scale to {y_scale:.3f} (image height: {height}, max zone y: {max_y})")
else:
y_scale = 1.0
# Apply more aggressive adjustment if image and zones are very different in scale
if max_x > width * 5 or max_x < width / 5:
x_scale = width / max_x
print(f"Major X scale adjustment to {x_scale:.3f}")
if max_y > height * 5 or max_y < height / 5:
y_scale = height / max_y
print(f"Major Y scale adjustment to {y_scale:.3f}")
# Apply the scaling to all zones
if x_scale != 1.0 or y_scale != 1.0:
for zone in zones:
zone['x1'] = int(zone['x1'] * x_scale)
zone['y1'] = int(zone['y1'] * y_scale)
zone['x2'] = int(zone['x2'] * x_scale)
zone['y2'] = int(zone['y2'] * y_scale)
print(f"Applied auto-scaling: X={x_scale}, Y={y_scale}")
except Exception as e:
print(f"Error parsing DocTags: {e}")
return False
# Create debug image with zones
create_debug_image(image, zones, page_num, output_path)
return True
def process_all_pages(pdf_path, doctags_path, output_base, dpi=200, scale=1.0, scale_x=None, scale_y=None, adjust=False):
"""Process all pages of the PDF and create visualizations."""
# Ensure results folder exists
results_dir = ensure_results_folder()
# Get total page count
total_pages = count_pdf_pages(pdf_path)
if total_pages == 0:
print("Error: Could not determine the number of pages in the PDF.")
return False
print(f"Processing all {total_pages} pages of the PDF...")
# Process each page
for page_num in range(1, total_pages + 1):
print(f"\nProcessing page {page_num} of {total_pages}...")
# Generate output paths for this page
if output_base is None:
output_path = results_dir / f"visualization_page_{page_num}.png"
else:
output_path = Path(output_base).with_stem(f"{Path(output_base).stem}_page_{page_num}")
# Process the page
process_page(pdf_path, page_num, doctags_path, output_path, dpi, scale, scale_x, scale_y, adjust)
# Create visualization
create_visualization(image, zones, page_num, output_path)
return True
def main():
# Parse arguments
args = parse_arguments()
# If just counting pages
if args.page_count:
page_count = count_pdf_pages(args.pdf)
print(f"The PDF has {page_count} pages.")
return
# Check if files exist
if not os.path.exists(args.pdf):
print(f"Error: PDF file not found: {args.pdf}")
@ -401,34 +186,15 @@ def main():
print(f"Error: DocTags file not found: {args.doctags}")
return
# Process page(s)
if args.page == 0: # Special case: process all pages
process_all_pages(
args.pdf,
args.doctags,
args.output,
args.dpi,
args.scale,
args.scale_x,
args.scale_y,
args.adjust
)
else:
# Determine output path
results_dir = ensure_results_folder()
output_path = args.output if args.output else results_dir / f"visualization_page_{args.page}.png"
process_page(
args.pdf,
args.page,
args.doctags,
output_path,
args.dpi,
args.scale,
args.scale_x,
args.scale_y,
args.adjust
)
# Process the page
process_page(
args.pdf,
args.page,
args.doctags,
args.output,
args.dpi,
args.adjust
)
if __name__ == "__main__":
main()

213
backend/utils.py Normal file
View file

@ -0,0 +1,213 @@
#!/usr/bin/env python3
"""
Common utilities for DocTags processing
"""
import os
import subprocess
import logging
from pathlib import Path
from typing import Optional, Tuple, Dict, List
import pdf2image
from pdf2image.pdf2image import pdfinfo_from_path
logger = logging.getLogger(__name__)
# Configuration constants
DEFAULT_DPI = 200
DEFAULT_GRID_SIZE = 500
MAX_WIDTH = 1200
RESULTS_DIR_NAME = "results"
def get_project_root() -> Path:
"""Get the project root directory."""
# If running from backend/page_treatment/, go up to root
current_file = Path(__file__)
if current_file.parent.name == 'page_treatment':
return current_file.parent.parent.parent
elif current_file.parent.name == 'backend':
return current_file.parent.parent
else:
return Path.cwd()
def ensure_results_folder(custom_path: Optional[str] = None) -> Path:
"""Create and return the results folder path."""
if custom_path:
results_dir = Path(custom_path)
else:
results_dir = get_project_root() / RESULTS_DIR_NAME
if not results_dir.exists():
results_dir.mkdir(parents=True)
logger.info(f"Created results directory: {results_dir}")
return results_dir
def count_pdf_pages(pdf_path: str) -> int:
"""Count the number of pages in a PDF file."""
if not os.path.exists(pdf_path):
logger.error(f"PDF file not found: {pdf_path}")
return 0
try:
info = pdfinfo_from_path(pdf_path)
return info["Pages"]
except Exception as e:
logger.warning(f"pdfinfo failed: {e}, trying fallback method")
try:
# Fallback: convert first page to check
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=1, last_page=1)
if not images:
return 0
# Binary search for last page
low, high = 1, 1000
while low < high:
mid = (low + high + 1) // 2
try:
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=mid, last_page=mid)
if images:
low = mid
else:
high = mid - 1
except:
high = mid - 1
return low
except Exception as e2:
logger.error(f"Error counting PDF pages: {e2}")
return 0
def load_pdf_page(pdf_path: str, page_num: int = 1, dpi: int = DEFAULT_DPI) -> Optional[object]:
"""Load a specific page from PDF as an image."""
if not os.path.exists(pdf_path):
raise FileNotFoundError(f"PDF file not found: {pdf_path}")
logger.info(f"Converting PDF page {page_num} to image (DPI: {dpi})...")
try:
pdf_images = pdf2image.convert_from_path(
pdf_path,
dpi=dpi,
first_page=page_num,
last_page=page_num
)
if not pdf_images:
raise Exception(f"Could not extract page {page_num} from PDF")
return pdf_images[0]
except Exception as e:
raise Exception(f"Error converting PDF to image: {e}")
def normalize_coordinates(elements: List[Dict], image_width: int, image_height: int,
grid_size: int = DEFAULT_GRID_SIZE) -> List[Dict]:
"""
Normalize coordinates from DocTags grid to actual image dimensions.
Args:
elements: List of elements with x1, y1, x2, y2 coordinates
image_width: Width of the image in pixels
image_height: Height of the image in pixels
grid_size: The grid size used in DocTags (default 500)
Returns:
List of elements with normalized coordinates
"""
normalized = []
for element in elements:
new_element = element.copy()
new_element['x1'] = int(element['x1'] * image_width / grid_size)
new_element['y1'] = int(element['y1'] * image_height / grid_size)
new_element['x2'] = int(element['x2'] * image_width / grid_size)
new_element['y2'] = int(element['y2'] * image_height / grid_size)
normalized.append(new_element)
return normalized
def auto_adjust_coordinates(elements: List[Dict], image_width: int, image_height: int) -> List[Dict]:
"""
Automatically adjust coordinates based on image dimensions.
"""
if not elements:
return elements
# Find maximum coordinates
max_x = max([el['x2'] for el in elements])
max_y = max([el['y2'] for el in elements])
# Check if coordinates are in normalized grid (0-500 range)
if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
logger.info(f"Detected normalized coordinates (0-{DEFAULT_GRID_SIZE} grid)")
return normalize_coordinates(elements, image_width, image_height)
# Calculate scaling factors
x_scale = calculate_scale_factor(max_x, image_width)
y_scale = calculate_scale_factor(max_y, image_height)
# Apply scaling
adjusted = []
for el in elements:
adjusted_el = el.copy()
adjusted_el['x1'] = int(el['x1'] * x_scale)
adjusted_el['y1'] = int(el['y1'] * y_scale)
adjusted_el['x2'] = int(el['x2'] * x_scale)
adjusted_el['y2'] = int(el['y2'] * y_scale)
adjusted.append(adjusted_el)
logger.info(f"Applied auto-scaling: X={x_scale:.3f}, Y={y_scale:.3f}")
return adjusted
def calculate_scale_factor(max_coord: float, image_size: float) -> float:
"""Calculate appropriate scaling factor."""
if max_coord <= 0:
return 1.0
# If coordinates are way off, apply aggressive scaling
if max_coord > image_size * 5 or max_coord < image_size / 5:
return image_size / max_coord
# Otherwise, apply conservative scaling
if max_coord > image_size:
return min(image_size / max_coord, 1.0)
else:
return max(image_size / max_coord, 0.5)
def run_command_with_timeout(command: str, timeout: int = 300, input_text: str = "n\n") -> Tuple[bool, str, str]:
"""
Run a command with timeout and return success, stdout, stderr.
"""
try:
process = subprocess.Popen(
command,
shell=True,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
universal_newlines=True
)
stdout, stderr = process.communicate(input=input_text, timeout=timeout)
success = process.returncode == 0
return success, stdout, stderr
except subprocess.TimeoutExpired:
process.kill()
return False, "", "Command timed out"
except Exception as e:
return False, "", str(e)
def format_duration(seconds: float) -> str:
"""Format duration in seconds to human readable format."""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
if hours > 0:
return f"{hours}:{minutes:02d}:{secs:02d}"
else:
return f"{minutes}:{secs:02d}"
def validate_coordinates(x1: int, y1: int, x2: int, y2: int,
width: int, height: int) -> bool:
"""Validate that coordinates are within bounds."""
return (0 <= x1 < x2 <= width and
0 <= y1 < y2 <= height)

860
frontend/static/app.js
View file

@ -1,107 +1,127 @@
// Keep track of active tasks
const activeTasks = {};
let pollingInterval = null;
// Keep track of generated outputs
const generatedOutputs = {
visualizer: null,
extractor: false
// DocTags Application State Management
const appState = {
activeTasks: {},
pollingInterval: null,
generatedOutputs: {
visualizer: null,
extractor: false
}
};
// Load and display PDF preview
// API Client
const api = {
async get(url) {
const response = await fetch(url);
if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`);
return response.json();
},
async post(url, formData) {
const response = await fetch(url, { method: 'POST', body: formData });
if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`);
return response.json();
}
};
// UI Helper Functions
const ui = {
show(elementId) {
document.getElementById(elementId).classList.remove('hidden');
},
hide(elementId) {
document.getElementById(elementId).classList.add('hidden');
},
setText(elementId, text) {
document.getElementById(elementId).textContent = text;
},
setHtml(elementId, html) {
document.getElementById(elementId).innerHTML = html;
},
getValue(elementId) {
return document.getElementById(elementId).value;
},
setValue(elementId, value) {
document.getElementById(elementId).value = value;
},
disable(elementId) {
document.getElementById(elementId).disabled = true;
},
enable(elementId) {
document.getElementById(elementId).disabled = false;
}
};
// PDF Preview Functions
function loadPDFPreview() {
const pdfFile = document.getElementById('pdf_file').value;
const pageNum = document.getElementById('page_num').value;
const pdfFile = ui.getValue('pdf_file');
const pageNum = ui.getValue('page_num');
if (!pdfFile) {
document.getElementById('pdf-preview-container').classList.add('hidden');
ui.hide('pdf-preview-container');
return;
}
// Update page number display
document.getElementById('preview-page-num').textContent = pageNum;
ui.setText('preview-page-num', pageNum);
ui.setValue('preview-page-input', pageNum);
// Update preview page input
const previewPageInput = document.getElementById('preview-page-input');
if (previewPageInput) {
previewPageInput.value = pageNum;
}
// Load the preview image
const previewImg = document.getElementById('pdf-preview-image');
previewImg.src = `/pdf-preview/${encodeURIComponent(pdfFile)}/${pageNum}`;
ui.show('pdf-preview-container');
// Show the preview container
document.getElementById('pdf-preview-container').classList.remove('hidden');
// Handle loading errors
previewImg.onerror = function() {
this.alt = 'Failed to load PDF preview';
console.error('Failed to load PDF preview');
};
}
// Navigate to a specific page in the preview
function navigateToPage(pageNum) {
const pdfFile = document.getElementById('pdf_file').value;
if (!pdfFile) return;
// Update the main page number input
document.getElementById('page_num').value = pageNum;
// Reload the preview
function changePreviewPage(delta) {
const currentPage = parseInt(ui.getValue('page_num'));
const newPage = Math.max(1, currentPage + delta);
ui.setValue('page_num', newPage);
loadPDFPreview();
}
// Handle preview page navigation
function changePreviewPage(delta) {
const currentPage = parseInt(document.getElementById('page_num').value);
const newPage = Math.max(1, currentPage + delta);
navigateToPage(newPage);
}
// Handle direct page input in preview
function goToPreviewPage() {
const pageInput = document.getElementById('preview-page-input');
const pageNum = parseInt(pageInput.value);
if (pageNum && pageNum > 0) {
navigateToPage(pageNum);
ui.setValue('page_num', pageNum);
loadPDFPreview();
} else {
// Reset to current page if invalid
pageInput.value = document.getElementById('page_num').value;
pageInput.value = ui.getValue('page_num');
}
}
// Tab switching functionality
// Tab Management
function switchTab(tabName) {
// Hide all tab contents
const tabContents = document.querySelectorAll('.tab-content');
tabContents.forEach(content => content.classList.remove('active'));
// Hide all tab contents and deactivate buttons
document.querySelectorAll('.tab-content').forEach(content => content.classList.remove('active'));
document.querySelectorAll('.tab-button').forEach(button => button.classList.remove('active'));
// Remove active class from all tab buttons
const tabButtons = document.querySelectorAll('.tab-button');
tabButtons.forEach(button => button.classList.remove('active'));
// Show selected tab content
// Show selected tab
document.getElementById(tabName + '-tab').classList.add('active');
// Add active class to selected tab button
event.target.classList.add('active');
// Show previously generated content when switching tabs
// Show previously generated content
if (tabName === 'analyzer') {
// Load PDF preview if a PDF is selected
loadPDFPreview();
} else if (tabName === 'visualizer' && generatedOutputs.visualizer) {
document.getElementById('result-image').src = generatedOutputs.visualizer + '?t=' + new Date().getTime();
document.getElementById('image-container').classList.remove('hidden');
} else if (tabName === 'extractor' && generatedOutputs.extractor) {
} else if (tabName === 'visualizer' && appState.generatedOutputs.visualizer) {
document.getElementById('result-image').src = appState.generatedOutputs.visualizer + '?t=' + Date.now();
ui.show('image-container');
} else if (tabName === 'extractor' && appState.generatedOutputs.extractor) {
loadExtractedImages();
}
}
// Update progress indicator
// Progress Management
function updateProgress(completedSteps) {
for (let i = 1; i <= 3; i++) {
const stepIndicator = document.getElementById(`step-${i}`);
@ -121,72 +141,202 @@ function updateProgress(completedSteps) {
}
}
// Load extracted images from the results folder
function loadExtractedImages() {
const imageGallery = document.getElementById('extracted-images-gallery');
const imageContainer = document.getElementById('extracted-images-container');
// First try to load the index.html to get the list of images
fetch('/results/pictures/index.html')
.then(response => {
if (!response.ok) {
throw new Error('No extracted images found');
}
return response.text();
})
.then(html => {
// Parse the HTML to extract image information
const parser = new DOMParser();
const doc = parser.parseFromString(html, 'text/html');
const imageCards = doc.querySelectorAll('.picture-card');
if (imageCards.length === 0) {
imageGallery.innerHTML = '<div class="no-images">No images were extracted from this page.</div>';
return;
}
let galleryHTML = '';
imageCards.forEach((card, index) => {
const img = card.querySelector('img');
const caption = card.querySelector('.picture-caption');
const coords = card.querySelector('.picture-coords');
if (img) {
const imgSrc = img.getAttribute('src');
const pictureId = index + 1;
const captionText = caption ? caption.textContent : '';
const coordsText = coords ? coords.textContent : '';
galleryHTML += `
<div class="extracted-image-card">
<div class="image-wrapper">
<img src="/results/pictures/${imgSrc}" alt="Extracted Image ${pictureId}"
onclick="openImageModal('/results/pictures/${imgSrc}', '${captionText}', '${coordsText}')">
</div>
<div class="image-info">
<h4>Picture ${pictureId}</h4>
${captionText ? `<p class="image-caption">${captionText}</p>` : ''}
<p class="image-coords">${coordsText}</p>
</div>
</div>
`;
}
});
imageGallery.innerHTML = galleryHTML;
imageContainer.classList.remove('hidden');
generatedOutputs.extractor = true;
})
.catch(error => {
console.log('No extracted images found:', error);
imageGallery.innerHTML = '<div class="no-images">No images have been extracted yet. Run the image extraction first.</div>';
generatedOutputs.extractor = false;
});
// Task Management
function startPolling() {
if (!appState.pollingInterval) {
appState.pollingInterval = setInterval(pollTasks, 1000);
}
}
// Open image in modal for better viewing
function stopPolling() {
if (Object.keys(appState.activeTasks).length === 0 && appState.pollingInterval) {
clearInterval(appState.pollingInterval);
appState.pollingInterval = null;
}
}
async function pollTasks() {
for (const taskId in appState.activeTasks) {
try {
const data = await api.get(`/task-status/${taskId}`);
updateTaskStatus(taskId, data);
} catch (error) {
console.error('Error polling task:', error);
}
}
}
function updateTaskStatus(taskId, data) {
const taskInfo = appState.activeTasks[taskId];
const statusElement = document.getElementById(`${taskInfo.type}-status`);
if (data.done) {
if (data.success) {
handleTaskSuccess(taskId, taskInfo, data, statusElement);
} else {
handleTaskFailure(taskId, data, statusElement);
}
delete appState.activeTasks[taskId];
stopPolling();
} else {
ui.setHtml(statusElement, '<div class="loader"></div><span class="working">Running...</span>');
ui.show(statusElement.id);
}
}
function handleTaskSuccess(taskId, taskInfo, data, statusElement) {
ui.setHtml(statusElement, '<span class="success">✓ Completed successfully!</span>');
ui.show(statusElement.id);
ui.enable(`${taskInfo.type}-btn`);
// Display output
ui.setText('output', data.output);
ui.show('output');
// Handle specific task types
if (taskInfo.type === 'visualizer' && data.image_file) {
appState.generatedOutputs.visualizer = '/' + data.image_file;
document.getElementById('result-image').src = appState.generatedOutputs.visualizer + '?t=' + Date.now();
ui.show('image-container');
} else if (taskInfo.type === 'extractor') {
setTimeout(loadExtractedImages, 1000);
}
// Update progress
const progressMap = { 'analyzer': 1, 'visualizer': 2, 'extractor': 3 };
updateProgress(progressMap[taskInfo.type]);
// Enable next steps
if (taskInfo.type === 'analyzer') {
ui.enable('visualizer-btn');
} else if (taskInfo.type === 'visualizer') {
ui.enable('extractor-btn');
}
}
function handleTaskFailure(taskId, data, statusElement) {
ui.setHtml(statusElement, '<span class="error">✗ Failed: ' + (data.error || 'Unknown error') + '</span>');
ui.show(statusElement.id);
ui.enable(statusElement.id.replace('-status', '-btn'));
ui.setText('output', 'Error: ' + (data.error || 'Unknown error'));
ui.show('output');
}
// Script Execution
async function runScript(script) {
const pdfFile = ui.getValue('pdf_file');
const pageNum = ui.getValue('page_num');
const adjust = document.getElementById('adjust').checked;
if (!pdfFile) {
alert('Please select a PDF file');
return;
}
// Disable button and show status
ui.disable(`${script}-btn`);
ui.setHtml(`${script}-status`, '<div class="loader"></div><span class="working">Starting...</span>');
ui.show(`${script}-status`);
ui.hide('output');
// Reset outputs for analyzer
if (script === 'analyzer') {
ui.hide('image-container');
ui.hide('extracted-images-container');
appState.generatedOutputs = { visualizer: null, extractor: false };
ui.disable('visualizer-btn');
ui.disable('extractor-btn');
updateProgress(0);
}
// Create form data
const formData = new FormData();
formData.append('pdf_file', pdfFile);
formData.append('page_num', pageNum);
formData.append('adjust', adjust);
try {
const data = await api.post(`/run-${script}`, formData);
if (data.success && data.task_id) {
appState.activeTasks[data.task_id] = {
type: script,
pageNum: pageNum
};
startPolling();
ui.setText('output', data.message || 'Task started, please wait...');
ui.show('output');
} else {
throw new Error(data.error || 'Failed to start task');
}
} catch (error) {
ui.setHtml(`${script}-status`, '<span class="error">✗ ' + error.message + '</span>');
ui.enable(`${script}-btn`);
ui.setText('output', 'Error: ' + error.message);
ui.show('output');
}
}
// Image Gallery Functions
async function loadExtractedImages() {
const imageGallery = document.getElementById('extracted-images-gallery');
try {
const response = await fetch('/results/pictures/index.html');
if (!response.ok) throw new Error('No extracted images found');
const html = await response.text();
const parser = new DOMParser();
const doc = parser.parseFromString(html, 'text/html');
const imageCards = doc.querySelectorAll('.picture-card');
if (imageCards.length === 0) {
imageGallery.innerHTML = '<div class="no-images">No images were extracted from this page.</div>';
return;
}
let galleryHTML = '';
imageCards.forEach((card, index) => {
const img = card.querySelector('img');
const caption = card.querySelector('.picture-caption');
const coords = card.querySelector('.picture-coords');
if (img) {
const imgSrc = img.getAttribute('src');
const pictureId = index + 1;
const captionText = caption ? caption.textContent : '';
const coordsText = coords ? coords.textContent : '';
galleryHTML += `
<div class="extracted-image-card">
<div class="image-wrapper">
<img src="/results/pictures/${imgSrc}" alt="Extracted Image ${pictureId}"
onclick="openImageModal('/results/pictures/${imgSrc}', '${captionText}', '${coordsText}')">
</div>
<div class="image-info">
<h4>Picture ${pictureId}</h4>
${captionText ? `<p class="image-caption">${captionText}</p>` : ''}
<p class="image-coords">${coordsText}</p>
</div>
</div>
`;
}
});
imageGallery.innerHTML = galleryHTML;
ui.show('extracted-images-container');
appState.generatedOutputs.extractor = true;
} catch (error) {
console.log('No extracted images found:', error);
imageGallery.innerHTML = '<div class="no-images">No images have been extracted yet. Run the image extraction first.</div>';
appState.generatedOutputs.extractor = false;
}
}
// Modal Functions
function openImageModal(imageSrc, caption, coords) {
// Create modal if it doesn't exist
let modal = document.getElementById('image-modal');
if (!modal) {
modal = document.createElement('div');
@ -205,403 +355,123 @@ function openImageModal(imageSrc, caption, coords) {
document.body.appendChild(modal);
}
// Set image and info
document.getElementById('modal-image').src = imageSrc;
document.getElementById('modal-caption').textContent = caption;
document.getElementById('modal-coords').textContent = coords;
// Show modal
ui.setText('modal-caption', caption);
ui.setText('modal-coords', coords);
modal.classList.add('active');
}
// Close image modal
function closeImageModal() {
const modal = document.getElementById('image-modal');
if (modal) {
modal.classList.remove('active');
}
if (modal) modal.classList.remove('active');
}
// Load PDF files on page load
window.addEventListener('DOMContentLoaded', function() {
const pdfStatus = document.getElementById('pdf-load-status');
pdfStatus.innerHTML = '<div class="loader"></div> Loading PDF files...';
fetch('/pdf-files')
.then(response => {
if (!response.ok) {
throw new Error('Failed to load PDF files');
}
return response.json();
})
.then(data => {
const select = document.getElementById('pdf_file');
if (data.length === 0) {
pdfStatus.innerHTML = '<span class="error">No PDF files found in the current directory</span>';
return;
}
data.forEach(file => {
const option = document.createElement('option');
option.value = file;
option.textContent = file;
select.appendChild(option);
});
pdfStatus.innerHTML = '<span class="success">Loaded ' + data.length + ' PDF files</span>';
})
.catch(error => {
console.error('Error loading PDFs:', error);
pdfStatus.innerHTML = '<span class="error">Error: ' + error.message + '</span>';
});
// Add event listeners for PDF selection and page number changes
document.getElementById('pdf_file').addEventListener('change', loadPDFPreview);
document.getElementById('page_num').addEventListener('change', loadPDFPreview);
// Add event listener for preview page input (if it exists)
const previewPageInput = document.getElementById('preview-page-input');
if (previewPageInput) {
previewPageInput.addEventListener('keypress', function(e) {
if (e.key === 'Enter') {
goToPreviewPage();
}
});
}
// Run an environment check on startup
checkEnvironment();
// Try to load any existing extracted images
loadExtractedImages();
// Check if there's an existing visualization
fetch('/results/visualization_page_1.png')
.then(response => {
if (response.ok) {
generatedOutputs.visualizer = '/results/visualization_page_1.png';
}
})
.catch(() => {
// No existing visualization
});
});
// Check the environment
function checkEnvironment() {
// Environment Check
async function checkEnvironment() {
const envDiv = document.getElementById('environment-check');
const envDetails = document.getElementById('env-details');
envDiv.classList.remove('hidden');
envDetails.innerHTML = '<div class="loader"></div> Checking environment...';
ui.show('environment-check');
ui.setHtml('env-details', '<div class="loader"></div> Checking environment...');
fetch('/check-environment')
.then(response => response.json())
.then(data => {
let html = '<ul>';
try {
const data = await api.get('/check-environment');
// Current working directory
html += '<li>Working directory: <code>' + data.cwd + '</code></li>';
let html = '<ul>';
html += `<li>Working directory: <code>${data.cwd}</code></li>`;
html += `<li>Python version: <code>${data.python_version}</code></li>`;
// Python version
html += '<li>Python version: <code>' + data.python_version + '</code></li>';
if (data.missing_scripts.length === 0) {
html += '<li class="env-success">✓ All required scripts found</li>';
} else {
html += `<li class="env-error">✗ Missing scripts: <code>${data.missing_scripts.join(', ')}</code></li>`;
}
// Required scripts check
if (data.missing_scripts.length === 0) {
html += '<li class="env-success">✓ All required scripts found</li>';
} else {
html += '<li class="env-error">✗ Missing scripts: <code>' + data.missing_scripts.join(', ') + '</code></li>';
}
if (data.pdf_files.length > 0) {
html += `<li class="env-success">✓ Found ${data.pdf_files.length} PDF files</li>`;
} else {
html += '<li class="env-error">✗ No PDF files found in the working directory</li>';
}
// PDF files check
if (data.pdf_files.length > 0) {
html += '<li class="env-success">✓ Found ' + data.pdf_files.length + ' PDF files: <code>' + data.pdf_files.join(', ') + '</code></li>';
} else {
html += '<li class="env-error">✗ No PDF files found in the working directory</li>';
}
// Results directory check
if (data.results_dir_exists) {
html += '<li class="env-success">✓ Results directory exists</li>';
if (data.results_dir_writable) {
html += '<li class="env-success">✓ Results directory is writable</li>';
} else {
html += '<li class="env-error">✗ Results directory is not writable</li>';
}
// Show files in results directory
if (data.results_files && data.results_files.length > 0) {
html += '<li>Files in results directory: <code>' + data.results_files.join(', ') + '</code></li>';
}
} else {
html += '<li class="env-error">✗ Results directory does not exist</li>';
}
html += '</ul>';
// List of all files for debugging
html += '<details><summary>All files in directory (' + data.files.length + ' files)</summary><pre>' +
data.files.join('\n') + '</pre></details>';
envDetails.innerHTML = html;
// Also check debug-results endpoint
fetch('/debug-results')
.then(response => response.json())
.then(debugData => {
html += '<details><summary>Results Directory Debug Info</summary><pre>' +
JSON.stringify(debugData, null, 2) + '</pre></details>';
envDetails.innerHTML = html;
})
.catch(error => {
console.error('Error getting debug info:', error);
});
})
.catch(error => {
envDetails.innerHTML = '<div class="env-error">Error checking environment: ' + error.message + '</div>';
});
html += '</ul>';
ui.setHtml('env-details', html);
} catch (error) {
ui.setHtml('env-details', `<div class="env-error">Error checking environment: ${error.message}</div>`);
}
}
// Manual command execution for debugging
function manuallyRunScript() {
const command = prompt("Enter command to run (e.g., 'python backend/page_treatment/analyzer.py --image document.pdf --page 1')");
// Manual Command Execution
async function manuallyRunScript() {
const command = prompt("Enter command to run:");
if (!command) return;
const outputDiv = document.getElementById('output');
outputDiv.textContent = 'Running command: ' + command + '\nPlease wait...';
outputDiv.classList.remove('hidden');
ui.setText('output', 'Running command: ' + command + '\nPlease wait...');
ui.show('output');
// Create form data
const formData = new FormData();
formData.append('command', command);
// Send the command directly to backend
fetch('/run-manual-command', {
method: 'POST',
body: formData
})
.then(response => response.json())
.then(data => {
outputDiv.textContent = 'Command: ' + command + '\n\n' +
(data.success ? 'Success!\n\n' : 'Failed!\n\n') +
(data.output || '') +
(data.error ? '\n\nError: ' + data.error : '');
})
.catch(error => {
outputDiv.textContent = 'Error running command: ' + error.message;
try {
const data = await api.post('/run-manual-command', formData);
ui.setText('output',
'Command: ' + command + '\n\n' +
(data.success ? 'Success!\n\n' : 'Failed!\n\n') +
(data.output || '') +
(data.error ? '\n\nError: ' + data.error : '')
);
} catch (error) {
ui.setText('output', 'Error running command: ' + error.message);
}
}
// Initialize Application
window.addEventListener('DOMContentLoaded', async function() {
// Load PDF files
const pdfStatus = document.getElementById('pdf-load-status');
ui.setHtml('pdf-load-status', '<div class="loader"></div> Loading PDF files...');
try {
const data = await api.get('/pdf-files');
const select = document.getElementById('pdf_file');
if (data.length === 0) {
ui.setHtml('pdf-load-status', '<span class="error">No PDF files found</span>');
return;
}
data.forEach(file => {
const option = document.createElement('option');
option.value = file;
option.textContent = file;
select.appendChild(option);
});
}
// Start polling for task updates
function startPolling() {
if (pollingInterval) {
return; // Already polling
ui.setHtml('pdf-load-status', `<span class="success">Loaded ${data.length} PDF files</span>`);
} catch (error) {
ui.setHtml('pdf-load-status', '<span class="error">Error: ' + error.message + '</span>');
}
pollingInterval = setInterval(pollTasks, 1000);
}
// Add event listeners
document.getElementById('pdf_file').addEventListener('change', loadPDFPreview);
document.getElementById('page_num').addEventListener('change', loadPDFPreview);
// Stop polling when no active tasks
function checkAndStopPolling() {
if (Object.keys(activeTasks).length === 0 && pollingInterval) {
clearInterval(pollingInterval);
pollingInterval = null;
}
}
// Poll for task updates
function pollTasks() {
for (const taskId in activeTasks) {
const taskInfo = activeTasks[taskId];
fetch(`/task-status/${taskId}`)
.then(response => {
if (!response.ok) {
throw new Error('Failed to check task status');
}
return response.json();
})
.then(data => {
// Update status display
const statusElement = document.getElementById(`${taskInfo.type}-status`);
if (data.done) {
if (data.success) {
// Task completed successfully
statusElement.innerHTML = '<span class="success">✓ Completed successfully!</span>';
statusElement.classList.remove('hidden');
// Enable button
document.getElementById(`${taskInfo.type}-btn`).disabled = false;
// Display output
const outputDiv = document.getElementById('output');
outputDiv.textContent = data.output;
outputDiv.classList.remove('hidden');
// Show image if available for visualizer
if (taskInfo.type === 'visualizer' && data.image_file) {
// Fix: Ensure we're using the correct path format
generatedOutputs.visualizer = '/' + data.image_file;
const imagePath = generatedOutputs.visualizer + '?t=' + new Date().getTime();
console.log('Loading visualization from:', imagePath);
document.getElementById('result-image').src = imagePath;
document.getElementById('image-container').classList.remove('hidden');
}
// Load extracted images if extractor completed
if (taskInfo.type === 'extractor') {
setTimeout(() => {
loadExtractedImages();
}, 1000); // Small delay to ensure files are written
}
// Enable next step button and update progress
if (taskInfo.type === 'analyzer') {
document.getElementById('visualizer-btn').disabled = false;
updateProgress(1);
} else if (taskInfo.type === 'visualizer') {
document.getElementById('extractor-btn').disabled = false;
updateProgress(2);
} else if (taskInfo.type === 'extractor') {
updateProgress(3);
}
// Remove from active tasks
delete activeTasks[taskId];
checkAndStopPolling();
} else {
// Task failed
statusElement.innerHTML = '<span class="error">✗ Failed: ' + (data.error || 'Unknown error') + '</span>';
statusElement.classList.remove('hidden');
document.getElementById(`${taskInfo.type}-btn`).disabled = false;
// Display error output
const outputDiv = document.getElementById('output');
outputDiv.textContent = 'Error: ' + (data.error || 'Unknown error');
outputDiv.classList.remove('hidden');
// Remove from active tasks
delete activeTasks[taskId];
checkAndStopPolling();
}
} else {
// Still running
statusElement.innerHTML = '<div class="loader"></div><span class="working">Running...</span>';
statusElement.classList.remove('hidden');
}
})
.catch(error => {
console.error('Error checking task status:', error);
});
}
}
function runScript(script) {
const pdfFile = document.getElementById('pdf_file').value;
const pageNum = document.getElementById('page_num').value;
const adjust = document.getElementById('adjust').checked;
if (!pdfFile) {
alert('Please select a PDF file');
return;
}
// Disable button
const button = document.getElementById(`${script}-btn`);
button.disabled = true;
// Create form data
const formData = new FormData();
formData.append('pdf_file', pdfFile);
formData.append('page_num', pageNum);
formData.append('adjust', adjust);
// Show running status
const statusElement = document.getElementById(`${script}-status`);
statusElement.innerHTML = '<div class="loader"></div><span class="working">Starting...</span>';
statusElement.classList.remove('hidden');
// Clear previous output
document.getElementById('output').classList.add('hidden');
// Don't hide content when switching between tabs
// Only hide when running a new analysis on the same tab
if (script === 'analyzer') {
// Reset everything when running analyzer
document.getElementById('image-container').classList.add('hidden');
document.getElementById('extracted-images-container').classList.add('hidden');
generatedOutputs.visualizer = null;
generatedOutputs.extractor = false;
}
// Determine endpoint
let endpoint;
switch(script) {
case 'analyzer':
endpoint = '/run-analyzer';
// Disable next step buttons
document.getElementById('visualizer-btn').disabled = true;
document.getElementById('extractor-btn').disabled = true;
updateProgress(0);
break;
case 'visualizer':
endpoint = '/run-visualizer';
// Disable next step button
document.getElementById('extractor-btn').disabled = true;
break;
case 'extractor':
endpoint = '/run-extractor';
break;
}
// Send request
fetch(endpoint, {
method: 'POST',
body: formData
})
.then(response => {
if (!response.ok) {
throw new Error('Failed to start task');
}
return response.json();
})
.then(data => {
if (data.success && data.task_id) {
// Store task information
activeTasks[data.task_id] = {
type: script,
pageNum: pageNum
};
// Start polling for updates
startPolling();
// Update status
statusElement.innerHTML = '<div class="loader"></div><span class="working">Running...</span>';
// Show initial output
const outputDiv = document.getElementById('output');
outputDiv.textContent = data.message || 'Task started, please wait...';
outputDiv.classList.remove('hidden');
} else {
// Failed to start task
statusElement.innerHTML = '<span class="error">✗ Failed to start task</span>';
button.disabled = false;
// Show error
const outputDiv = document.getElementById('output');
outputDiv.textContent = 'Error: ' + (data.error || 'Failed to start task');
outputDiv.classList.remove('hidden');
}
})
.catch(error => {
console.error('Error starting task:', error);
statusElement.innerHTML = '<span class="error">✗ ' + error.message + '</span>';
button.disabled = false;
// Show error
const outputDiv = document.getElementById('output');
outputDiv.textContent = 'Error: ' + error.message;
outputDiv.classList.remove('hidden');
const previewPageInput = document.getElementById('preview-page-input');
if (previewPageInput) {
previewPageInput.addEventListener('keypress', function(e) {
if (e.key === 'Enter') goToPreviewPage();
});
}
}
// Initial checks
checkEnvironment();
loadExtractedImages();
// Check for existing visualization
try {
const response = await fetch('/results/visualization_page_1.png');
if (response.ok) {
appState.generatedOutputs.visualizer = '/results/visualization_page_1.png';
}
} catch {}
});