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

314
backend/batch_treatment/batch_processor.py
View file

@ -1,7 +1,6 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
""" """
Batch Processor for DocTags Batch Processor for DocTags - Handles parallel processing of PDF documents
Handles batch processing of PDF documents with parallel processing support
""" """
import os import os
@ -14,18 +13,15 @@ import logging
from pathlib import Path from pathlib import Path
from datetime import datetime from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
import subprocess
import shutil import shutil
import zipfile 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__))))) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
# Configure logging from backend.utils import ensure_results_folder, run_command_with_timeout, format_duration
logging.basicConfig( from backend.config import BATCH_WORKERS, PROCESSING_TIMEOUT
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -47,7 +43,7 @@ class BatchProcessor:
'completed': False, 'completed': False,
'paused': False, 'paused': False,
'cancelled': False, 'cancelled': False,
'page_statuses': {}, 'page_statuses': {str(page): 'pending' for page in range(start_page, end_page + 1)},
'stages': { 'stages': {
'analysis': {'completed': 0, 'total': self.total_pages}, 'analysis': {'completed': 0, 'total': self.total_pages},
'visualization': {'completed': 0, 'total': self.total_pages}, 'visualization': {'completed': 0, 'total': self.total_pages},
@ -62,17 +58,13 @@ class BatchProcessor:
'logs': [] 'logs': []
} }
# Initialize page statuses
for page in range(start_page, end_page + 1):
self.state['page_statuses'][str(page)] = 'pending'
# Threading # Threading
self.lock = threading.Lock() self.lock = threading.Lock()
self.pause_event = threading.Event() self.pause_event = threading.Event()
self.pause_event.set() # Start unpaused self.pause_event.set() # Start unpaused
# Create batch results directory # 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) self.results_dir.mkdir(parents=True, exist_ok=True)
# Log file # Log file
@ -89,11 +81,11 @@ class BatchProcessor:
with self.lock: with self.lock:
self.state['logs'].append(log_entry) 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: if len(self.state['logs']) > 100:
self.state['logs'] = 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: with open(self.log_file, 'a') as f:
f.write(f"[{timestamp}] [{level.upper()}] {message}\n") f.write(f"[{timestamp}] [{level.upper()}] {message}\n")
@ -125,7 +117,7 @@ class BatchProcessor:
raise Exception("Analyzer failed") raise Exception("Analyzer failed")
self.update_stage_progress('analysis') self.update_stage_progress('analysis')
# Check if paused or cancelled # Check pause/cancel
self.pause_event.wait() self.pause_event.wait()
if self.state['cancelled']: if self.state['cancelled']:
return False return False
@ -135,7 +127,7 @@ class BatchProcessor:
raise Exception("Visualizer failed") raise Exception("Visualizer failed")
self.update_stage_progress('visualization') self.update_stage_progress('visualization')
# Check if paused or cancelled # Check pause/cancel
self.pause_event.wait() self.pause_event.wait()
if self.state['cancelled']: if self.state['cancelled']:
return False return False
@ -152,7 +144,6 @@ class BatchProcessor:
self.update_page_status(page_num, 'completed') self.update_page_status(page_num, 'completed')
self.log_message(f"Successfully processed page {page_num}", 'success') self.log_message(f"Successfully processed page {page_num}", 'success')
return True return True
except Exception as e: except Exception as e:
@ -167,45 +158,26 @@ class BatchProcessor:
self.update_page_status(page_num, 'failed') self.update_page_status(page_num, 'failed')
self.log_message(f"Failed to process page {page_num}: {str(e)}", 'error') self.log_message(f"Failed to process page {page_num}: {str(e)}", 'error')
return False return False
def run_analyzer(self, page_num): def run_analyzer(self, page_num):
"""Run the analyzer for a specific page""" """Run the analyzer for a specific page"""
try: try:
# Use standard results directory for analyzer output command = (f"python backend/page_treatment/analyzer.py "
output_base = Path("results") / f"output" f"--image {self.pdf_file} --page {page_num} "
f"--start-page {page_num} --end-page {page_num}")
# 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)
]
self.log_message(f"Running analyzer for page {page_num}") self.log_message(f"Running analyzer for page {page_num}")
# Run the command success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
process = subprocess.Popen(
command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True
)
# Send "n" to bypass prompts if not success:
stdout, stderr = process.communicate(input="n\n", timeout=300) # 5 minute timeout
if process.returncode != 0:
raise Exception(f"Analyzer failed: {stderr}") raise Exception(f"Analyzer failed: {stderr}")
# Copy doctags to batch directory for archiving # Copy doctags to batch directory
doctags_src = Path("results") / "output.doctags.txt" doctags_src = ensure_results_folder() / "output.doctags.txt"
doctags_dst = self.results_dir / f"page_{page_num}.doctags.txt" doctags_dst = self.results_dir / f"page_{page_num}.doctags.txt"
if doctags_src.exists(): if doctags_src.exists():
shutil.copy2(doctags_src, doctags_dst) shutil.copy2(doctags_src, doctags_dst)
self.log_message(f"DocTags saved for page {page_num}") self.log_message(f"DocTags saved for page {page_num}")
@ -221,47 +193,33 @@ class BatchProcessor:
def run_visualizer(self, page_num): def run_visualizer(self, page_num):
"""Run the visualizer for a specific page""" """Run the visualizer for a specific page"""
try: try:
# The visualizer expects doctags in the standard location # Ensure correct doctags file is in place
doctags_path = Path("results") / "output.doctags.txt" doctags_path = ensure_results_folder() / "output.doctags.txt"
# First, ensure we have the right doctags file for this page
page_doctags = self.results_dir / f"page_{page_num}.doctags.txt" page_doctags = self.results_dir / f"page_{page_num}.doctags.txt"
if page_doctags.exists(): if page_doctags.exists():
shutil.copy2(page_doctags, doctags_path) shutil.copy2(page_doctags, doctags_path)
# Update path to use backend directory command = (f"python backend/page_treatment/visualizer.py "
command = [ f"--doctags {doctags_path} --pdf {self.pdf_file} --page {page_num}")
"python", "backend/page_treatment/visualizer.py",
"--doctags", str(doctags_path),
"--pdf", self.pdf_file,
"--page", str(page_num)
]
if self.options.get('adjust', True): if self.options.get('adjust', True):
command.append("--adjust") command += " --adjust"
self.log_message(f"Running visualizer for page {page_num}") self.log_message(f"Running visualizer for page {page_num}")
# Run the command success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
process = subprocess.run(
command,
capture_output=True,
text=True,
timeout=300
)
if process.returncode != 0: if not success:
raise Exception(f"Visualizer failed: {process.stderr}") raise Exception(f"Visualizer failed: {stderr}")
# The visualizer should have created the file in results/ # Copy visualization to batch directory
viz_src = Path("results") / f"visualization_page_{page_num}.png" viz_src = ensure_results_folder() / f"visualization_page_{page_num}.png"
viz_dst = self.results_dir / f"visualization_page_{page_num}.png" viz_dst = self.results_dir / f"visualization_page_{page_num}.png"
if viz_src.exists(): if viz_src.exists():
shutil.copy2(viz_src, viz_dst) shutil.copy2(viz_src, viz_dst)
self.log_message(f"Visualization saved for page {page_num}") 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 return True
@ -272,39 +230,29 @@ class BatchProcessor:
def run_extractor(self, page_num): def run_extractor(self, page_num):
"""Run the picture extractor for a specific page""" """Run the picture extractor for a specific page"""
try: try:
# Ensure we have the right doctags file for this page # Ensure correct doctags file is in place
doctags_path = Path("results") / "output.doctags.txt" doctags_path = ensure_results_folder() / "output.doctags.txt"
page_doctags = self.results_dir / f"page_{page_num}.doctags.txt" page_doctags = self.results_dir / f"page_{page_num}.doctags.txt"
if page_doctags.exists(): if page_doctags.exists():
shutil.copy2(page_doctags, doctags_path) shutil.copy2(page_doctags, doctags_path)
# Update path to use backend directory command = (f"python backend/page_treatment/picture_extractor.py "
command = [ f"--doctags {doctags_path} --pdf {self.pdf_file} --page {page_num}")
"python", "backend/page_treatment/picture_extractor.py",
"--doctags", str(doctags_path),
"--pdf", self.pdf_file,
"--page", str(page_num)
]
if self.options.get('adjust', True): if self.options.get('adjust', True):
command.append("--adjust") command += " --adjust"
self.log_message(f"Running extractor for page {page_num}") self.log_message(f"Running extractor for page {page_num}")
# Run the command success, stdout, stderr = run_command_with_timeout(command, PROCESSING_TIMEOUT)
process = subprocess.run(
command,
capture_output=True,
text=True,
timeout=300
)
if process.returncode != 0: if not success:
self.log_message(f"Extractor warning for page {page_num}: {process.stderr}", 'warning') self.log_message(f"Extractor warning for page {page_num}: {stderr}", 'warning')
# Count and copy extracted images # Count and copy extracted images
image_count = 0 image_count = 0
pics_src = Path("results") / "pictures" pics_src = ensure_results_folder() / "pictures"
pics_dst = self.results_dir / f"pictures_page_{page_num}" pics_dst = self.results_dir / f"pictures_page_{page_num}"
if pics_src.exists(): if pics_src.exists():
@ -313,8 +261,8 @@ class BatchProcessor:
shutil.rmtree(pics_dst) shutil.rmtree(pics_dst)
shutil.copytree(pics_src, pics_dst) shutil.copytree(pics_src, pics_dst)
# Also copy to page-specific location for web interface # Also copy for web interface
pics_web = Path("results") / f"pictures_page_{page_num}" pics_web = ensure_results_folder() / f"pictures_page_{page_num}"
if pics_web.exists(): if pics_web.exists():
shutil.rmtree(pics_web) shutil.rmtree(pics_web)
shutil.copytree(pics_src, pics_web) shutil.copytree(pics_src, pics_web)
@ -322,8 +270,6 @@ class BatchProcessor:
# Count PNG files # Count PNG files
image_count = len(list(pics_dst.glob("*.png"))) image_count = len(list(pics_dst.glob("*.png")))
self.log_message(f"Extracted {image_count} images from page {page_num}") 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 return image_count
@ -335,10 +281,11 @@ class BatchProcessor:
"""Main batch processing loop""" """Main batch processing loop"""
try: try:
self.state['status'] = 'processing' 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 # 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 # Create page list
pages = list(range(self.start_page, self.end_page + 1)) pages = list(range(self.start_page, self.end_page + 1))
@ -347,7 +294,8 @@ class BatchProcessor:
if max_workers > 1: if max_workers > 1:
# Parallel processing # Parallel processing
with ThreadPoolExecutor(max_workers=max_workers) as executor: 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): for future in as_completed(futures):
if self.state['cancelled']: if self.state['cancelled']:
@ -358,7 +306,8 @@ class BatchProcessor:
try: try:
future.result() future.result()
except Exception as e: 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: else:
# Sequential processing # Sequential processing
for page in pages: for page in pages:
@ -376,7 +325,8 @@ class BatchProcessor:
self.state['status'] = 'cancelled' if self.state['cancelled'] else 'completed' self.state['status'] = 'cancelled' if self.state['cancelled'] else 'completed'
duration = time.time() - self.state['start_time'] 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: except Exception as e:
self.log_message(f"Critical error in batch processing: {str(e)}", 'error') self.log_message(f"Critical error in batch processing: {str(e)}", 'error')
@ -385,157 +335,100 @@ class BatchProcessor:
self.state['status'] = 'error' self.state['status'] = 'error'
def generate_report(self): def generate_report(self):
"""Generate a comprehensive HTML report""" """Generate HTML report of batch processing results"""
try: try:
self.log_message("Generating batch processing report") self.log_message("Generating batch processing report")
report_path = self.results_dir / "report.html"
# Calculate statistics
duration = time.time() - self.state['start_time'] 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> <html>
<head> <head>
<meta charset="UTF-8"> <meta charset="UTF-8">
<title>Batch Processing Report - {self.pdf_file}</title> <title>Batch Processing Report - {self.pdf_file}</title>
<style> <style>
body {{ font-family: Arial, sans-serif; margin: 40px; background: #f5f5f5; }} 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); }} .container {{ max-width: 1200px; margin: 0 auto; background: white;
h1 {{ color: #2c3e50; margin-bottom: 10px; }} padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }}
.subtitle {{ color: #7f8c8d; margin-bottom: 30px; }} h1 {{ color: #2c3e50; }}
.stats {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 20px; margin: 30px 0; }} .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-box {{ background: #ecf0f1; padding: 20px; border-radius: 8px; text-align: center; }}
.stat-value {{ font-size: 2.5em; font-weight: bold; color: #2c3e50; }} .stat-value {{ font-size: 2.5em; font-weight: bold; color: #2c3e50; }}
.stat-label {{ color: #7f8c8d; margin-top: 5px; }}
.success {{ color: #27ae60; }} .success {{ color: #27ae60; }}
.error {{ color: #e74c3c; }} .error {{ color: #e74c3c; }}
table {{ width: 100%; border-collapse: collapse; margin: 20px 0; }} table {{ width: 100%; border-collapse: collapse; margin: 20px 0; }}
th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ecf0f1; }} th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ecf0f1; }}
th {{ background: #34495e; color: white; }} 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> </style>
</head> </head>
<body> <body>
<div class="container"> <div class="container">
<h1>Batch Processing Report</h1> <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="stats">
<div class="stat-box"> <div class="stat-box">
<div class="stat-value">{self.total_pages}</div> <div class="stat-value">{self.total_pages}</div>
<div class="stat-label">Total Pages</div> <div>Total Pages</div>
</div> </div>
<div class="stat-box"> <div class="stat-box">
<div class="stat-value success">{self.state['results']['successful']}</div> <div class="stat-value success">{self.state['results']['successful']}</div>
<div class="stat-label">Successful</div> <div>Successful</div>
</div> </div>
<div class="stat-box"> <div class="stat-box">
<div class="stat-value error">{self.state['results']['failed']}</div> <div class="stat-value error">{self.state['results']['failed']}</div>
<div class="stat-label">Failed</div> <div>Failed</div>
</div> </div>
<div class="stat-box"> <div class="stat-box">
<div class="stat-value">{self.state['results']['totalImages']}</div> <div class="stat-value">{self.state['results']['totalImages']}</div>
<div class="stat-label">Images Extracted</div> <div>Images Extracted</div>
</div> </div>
<div class="stat-box"> <div class="stat-box">
<div class="stat-value">{self.format_duration(duration)}</div> <div class="stat-value">{format_duration(duration)}</div>
<div class="stat-label">Processing Time</div> <div>Processing Time</div>
</div> </div>
<div class="stat-box"> <div class="stat-box">
<div class="stat-value">{success_rate:.1f}%</div> <div class="stat-value">{success_rate:.1f}%</div>
<div class="stat-label">Success Rate</div> <div>Success Rate</div>
</div> </div>
</div> </div>
"""
<h2>Processing Details</h2> # Add failed pages if any
if self.state['results']['failedPages']:
html += """
<h2>Failed Pages</h2>
<table> <table>
<tr> <tr><th>Page Number</th><th>Reason</th></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>
""" """
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 html += """
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>
</div> </div>
</body> </body>
</html> </html>
""" """
return 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')
def pause(self): def pause(self):
"""Pause the batch processing""" """Pause the batch processing"""
@ -552,11 +445,11 @@ class BatchProcessor:
def cancel(self): def cancel(self):
"""Cancel the batch processing""" """Cancel the batch processing"""
self.state['cancelled'] = True self.state['cancelled'] = True
self.pause_event.set() # Ensure not stuck on pause self.pause_event.set()
self.log_message("Batch processing cancelled") self.log_message("Batch processing cancelled")
def get_state(self): def get_state(self):
"""Get the current state with calculated fields""" """Get current state with calculated fields"""
with self.lock: with self.lock:
state = self.state.copy() state = self.state.copy()
@ -572,24 +465,12 @@ class BatchProcessor:
return state 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): def create_zip_archive(self):
"""Create a ZIP archive of all results""" """Create ZIP archive of all results"""
try: try:
zip_path = self.results_dir / f"batch_results_{self.batch_id}.zip" zip_path = self.results_dir / f"batch_results_{self.batch_id}.zip"
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: 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('*'): for file_path in self.results_dir.rglob('*'):
if file_path.is_file() and file_path != zip_path: if file_path.is_file() and file_path != zip_path:
arcname = file_path.relative_to(self.results_dir) arcname = file_path.relative_to(self.results_dir)
@ -649,3 +530,4 @@ def cleanup_old_batches(max_age_hours=24):
for batch_id in to_remove: 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", # "requests",
# "argparse", # "argparse",
# "pdf2image", # "pdf2image",
# "pymupdf", # Optional for better PDF handling
# ] # ]
# /// # ///
import argparse import argparse
import os import os
import tempfile import tempfile
import re import re
import base64
from io import BytesIO
from pathlib import Path from pathlib import Path
from urllib.parse import urlparse from urllib.parse import urlparse
import requests import requests
from PIL import Image, UnidentifiedImageError from PIL import Image
from pdf2image import convert_from_path, convert_from_bytes from pdf2image import convert_from_bytes
from docling_core.types.doc import ImageRefMode from docling_core.types.doc import ImageRefMode
from docling_core.types.doc.document import DocTagsDocument, DoclingDocument 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(): from backend.utils import ensure_results_folder, load_pdf_page, get_project_root
"""Create the results folder if it doesn't exist.""" from backend.config import MODEL_PATH, MAX_TOKENS, DEFAULT_DPI
# 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(): def parse_arguments():
"""Parse command line arguments.""" """Parse command line arguments."""
@ -53,257 +41,69 @@ def parse_arguments():
help='Prompt for the model') help='Prompt for the model')
parser.add_argument('--output', '-o', type=str, default=str(results_dir / "output.html"), parser.add_argument('--output', '-o', type=str, default=str(results_dir / "output.html"),
help='Output file path') 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, parser.add_argument('--page', type=int, default=1,
help='Page number to process for PDF files (starts at 1)') help='Page number to process for PDF files (starts at 1)')
parser.add_argument('--dpi', type=int, default=200, parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering (higher values produce larger images)') help='DPI for PDF rendering')
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('--start-page', type=int, default=1, parser.add_argument('--start-page', type=int, default=1,
help='Start processing PDF from this page number') help='Start processing PDF from this page number')
parser.add_argument('--end-page', type=int, default=None, parser.add_argument('--end-page', type=int, default=None,
help='Stop processing PDF at this page number') 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() return parser.parse_args()
def load_image(image_path, page_num=1, dpi=DEFAULT_DPI):
def load_image(image_path, page_num=1, dpi=200):
"""Load image from URL, local image file, or PDF.""" """Load image from URL, local image file, or PDF."""
if urlparse(image_path).scheme in ['http', 'https']: # it is a URL if urlparse(image_path).scheme in ['http', 'https']:
try: response = requests.get(image_path, stream=True, timeout=10)
response = requests.get(image_path, stream=True, timeout=10) response.raise_for_status()
response.raise_for_status()
content = response.content
# Check if it's a PDF if image_path.lower().endswith('.pdf') or response.headers.get('Content-Type') == 'application/pdf':
if image_path.lower().endswith('.pdf') or response.headers.get('Content-Type') == 'application/pdf': print(f"Converting PDF from URL (page {page_num})...")
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)
pdf_images = convert_from_bytes(content, dpi=dpi, first_page=page_num, last_page=page_num) if not pdf_images:
if not pdf_images: raise Exception(f"Could not extract page {page_num} from PDF")
raise Exception(f"Could not extract page {page_num} from PDF") return pdf_images[0]
return pdf_images[0] # Return the first (and only) page else:
else: return Image.open(response.raw)
return Image.open(BytesIO(content)) else:
except requests.exceptions.RequestException as e:
raise Exception(f"Error loading image from URL: {e}")
else: # it is a local file
image_path = Path(image_path) image_path = Path(image_path)
if not image_path.exists(): if not image_path.exists():
raise FileNotFoundError(f"File not found: {image_path}") raise FileNotFoundError(f"File not found: {image_path}")
# Check if it's a PDF
if image_path.suffix.lower() == '.pdf': if image_path.suffix.lower() == '.pdf':
print(f"Converting PDF to image (page {page_num}, DPI: {dpi})...") return load_pdf_page(str(image_path), page_num, 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}")
else: else:
try: return Image.open(image_path)
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.")
def process_page(model, processor, config, args, pil_image, page_num=1):
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):
"""Process a single page from a PDF or image file.""" """Process a single page from a PDF or image file."""
from mlx_vlm.prompt_utils import apply_chat_template from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import stream_generate from mlx_vlm.utils import stream_generate
# Ensure results folder exists
results_dir = ensure_results_folder() results_dir = ensure_results_folder()
# Prepare input
prompt = args.prompt
output_base = Path(args.output) output_base = Path(args.output)
output_path = output_base
# If processing PDF and output is a path without explicit numbering, add page numbers # Handle multi-page output naming
if Path(image_path).suffix.lower() == '.pdf' and page_num > 1: if Path(args.image).suffix.lower() == '.pdf' and page_num > 1:
# Get base filename without extension output_path = results_dir / f"{output_base.stem}_page{page_num}{output_base.suffix}"
base_name = output_base.stem else:
output_path = results_dir / f"{base_name}_page{page_num}{output_base.suffix}" output_path = output_base
print(f"Processing page {page_num}, output will be saved to {output_path}") 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: with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img_file:
temp_img_path = temp_img_file.name temp_img_path = temp_img_file.name
pil_image.save(temp_img_path, format='PNG') pil_image.save(temp_img_path, format='PNG')
print(f"Saved temporary image to: {temp_img_path}")
try: try:
# Apply chat template # Apply chat template and generate
formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1) formatted_prompt = apply_chat_template(processor, config, args.prompt, num_images=1)
# Generate output
print(f"Generating DocTags for page {page_num}: \n\n") print(f"Generating DocTags for page {page_num}: \n\n")
output = "" output = ""
for token in stream_generate( 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 output += token.text
print(token.text, end="") print(token.text, end="")
@ -311,184 +111,52 @@ def process_page(args, model, processor, config, image_path, pil_image, page_num
break break
print("\n\n") 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: finally:
# Clean up the temporary file # Clean up temporary file
if os.path.exists(temp_img_path): if os.path.exists(temp_img_path):
os.unlink(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" doctags_path = results_dir / f"{output_path.stem}.doctags.txt"
with open(doctags_path, 'w', encoding='utf-8') as f: with open(doctags_path, 'w', encoding='utf-8') as f:
f.write(output) f.write(output)
print(f"Raw DocTags saved to: {doctags_path}") 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 return output_path
def main(): def main():
# Ensure results folder exists
ensure_results_folder()
# Parse arguments
args = parse_arguments() args = parse_arguments()
# Settings
DEBUG_MODE = args.debug
DOCTAGS_ONLY = args.doctags_only
image_path = args.image
# Load the model # Load the model
print("Loading model...") print("Loading model...")
try: try:
from mlx_vlm import load, generate from mlx_vlm import load
from mlx_vlm.utils import load_config from mlx_vlm.utils import load_config
model_path = "ds4sd/SmolDocling-256M-preview-mlx-bf16" model, processor = load(MODEL_PATH)
model, processor = load(model_path) config = load_config(MODEL_PATH)
config = load_config(model_path)
except Exception as e: except Exception as e:
print(f"Error loading model: {e}") print(f"Error loading model: {e}")
import traceback
traceback.print_exc()
return return
# Check if the input is a PDF # Process the image/PDF
pdf_path = Path(image_path) try:
is_pdf = False # Handle single page or range
pdf_page_count = 0 start_page = args.start_page
end_page = args.end_page or args.page
if pdf_path.suffix.lower() == '.pdf': for page_num in range(start_page, end_page + 1):
is_pdf = True print(f"\nProcessing page {page_num}...")
# 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
# Determine which pages to process pil_image = load_image(args.image, page_num=page_num, dpi=args.dpi)
process_all_pages = args.all_pages print(f"Page {page_num} loaded: {pil_image.size}")
start_page = args.start_page
end_page = args.end_page if args.end_page else pdf_page_count
max_pages = args.max_pages
# Validate page ranges process_page(model, processor, config, args, pil_image, page_num)
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()
except Exception as e:
print(f"Error processing: {e}")
import traceback
traceback.print_exc()
if __name__ == "__main__": if __name__ == "__main__":
main() main()

259
backend/page_treatment/picture_extractor.py
View file

@ -1,41 +1,26 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
""" """
DocTags Picture Extractor - Extract <picture> elements from DocTags and save as separate image files. DocTags Picture Extractor - Extract <picture> elements from DocTags.
Usage:
python picture_extractor.py --doctags output.doctags.txt --pdf document.pdf --page 1 --output pictures
""" """
import argparse import argparse
import os import os
import re import re
import sys
from io import BytesIO
from pathlib import Path from pathlib import Path
import pdf2image
from PIL import Image 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 # Regular expression to extract picture location data
PICTURE_PATTERN = r'<picture>.*?<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>(.*?)</picture>' 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(): def parse_arguments():
"""Parse command line arguments.""" """Parse command line arguments."""
results_dir = ensure_results_folder() results_dir = ensure_results_folder()
@ -46,46 +31,19 @@ def parse_arguments():
parser.add_argument('--pdf', '-p', type=str, required=True, parser.add_argument('--pdf', '-p', type=str, required=True,
help='Path to original PDF file') help='Path to original PDF file')
parser.add_argument('--page', type=int, default=1, 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"), parser.add_argument('--output', '-o', type=str, default=str(results_dir / "pictures"),
help='Output directory for extracted pictures') help='Output directory for extracted pictures')
parser.add_argument('--dpi', type=int, default=300, parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering (higher values produce larger images)') help='DPI for PDF rendering')
parser.add_argument('--max-width', type=int, default=1200, parser.add_argument('--max-width', type=int, default=MAX_IMAGE_WIDTH,
help='Maximum width of output images in pixels') help='Maximum width of output images in pixels')
parser.add_argument('--adjust', action='store_true', parser.add_argument('--adjust', action='store_true',
help='Try to automatically adjust scaling') 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, parser.add_argument('--margin', type=int, default=0,
help='Add margin around extracted pictures in pixels') 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() 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): def extract_pictures_from_doctags(doctags_path):
"""Parse DocTags file and extract picture elements with their coordinates.""" """Parse DocTags file and extract picture elements with their coordinates."""
if not os.path.exists(doctags_path): 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: with open(doctags_path, 'r', encoding='utf-8') as f:
doctags_content = f.read() doctags_content = f.read()
# Find all picture elements with location information
pictures = [] pictures = []
picture_matches = re.finditer(PICTURE_PATTERN, doctags_content, re.DOTALL) picture_matches = re.finditer(PICTURE_PATTERN, doctags_content, re.DOTALL)
for i, match in enumerate(picture_matches): for i, match in enumerate(picture_matches):
x1, y1, x2, y2, caption = match.groups() 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() clean_caption = re.sub(r'<loc_\d+>', '', caption).strip()
pictures.append({ pictures.append({
'id': i + 1, 'id': i + 1,
'x1': int(x1), 'x1': int(x1), 'y1': int(y1),
'y1': int(y1), 'x2': int(x2), 'y2': int(y2),
'x2': int(x2),
'y2': int(y2),
'caption': clean_caption 'caption': clean_caption
}) })
return pictures return pictures
def normalize_coordinates(pictures, image_width, image_height, grid_size=500): def extract_and_save_pictures(image, pictures, output_dir, max_width, margin):
"""
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):
"""Extract picture regions from the image and save them as separate files.""" """Extract picture regions from the image and save them as separate files."""
# Ensure output directory exists
output_path = ensure_results_folder(output_dir) output_path = ensure_results_folder(output_dir)
saved_files = [] saved_files = []
# Process each picture
for picture in pictures: for picture in pictures:
try: try:
# Add margin to coordinates if specified # Add margin to coordinates
x1 = max(0, picture['x1'] - margin) x1 = max(0, picture['x1'] - margin)
y1 = max(0, picture['y1'] - margin) y1 = max(0, picture['y1'] - margin)
x2 = min(image.width, picture['x2'] + margin) x2 = min(image.width, picture['x2'] + margin)
y2 = min(image.height, picture['y2'] + margin) y2 = min(image.height, picture['y2'] + margin)
# Check if coordinates are valid # Validate coordinates
if x1 >= x2 or y1 >= y2 or x1 < 0 or y1 < 0 or x2 > image.width or y2 > image.height: if not validate_coordinates(x1, y1, x2, y2, image.width, image.height):
print(f"Warning: Invalid coordinates for picture {picture['id']}: ({x1},{y1})-({x2},{y2})") print(f"Warning: Invalid coordinates for picture {picture['id']}")
continue continue
# Crop the image # 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) cropped_img = cropped_img.resize((max_width, new_height), Image.LANCZOS)
# Generate filename # Generate filename
caption = picture['caption'] if picture['caption']:
if caption: safe_caption = re.sub(r'[^\w\s-]', '', picture['caption'])[:30].strip().replace(' ', '_').lower()
# Create a filename-safe version of the caption (first 30 chars)
safe_caption = re.sub(r'[^\w\s-]', '', caption)[:30].strip().replace(' ', '_').lower()
filename = f"picture_{picture['id']}_{safe_caption}.png" filename = f"picture_{picture['id']}_{safe_caption}.png"
else: else:
filename = f"picture_{picture['id']}.png" 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 output_file = output_path / filename
cropped_img.save(output_file, format="PNG") cropped_img.save(output_file, format="PNG")
# Create a text file with the caption if available # Save caption if available
if caption: if picture['caption']:
caption_file = output_path / f"{output_file.stem}.txt" caption_file = output_path / f"{output_file.stem}.txt"
with open(caption_file, 'w', encoding='utf-8') as f: 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}") print(f"Saved picture {picture['id']} to {output_file}")
saved_files.append(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 {{ .picture-info {{
padding: 15px; padding: 15px;
}} }}
.picture-caption {{
margin-top: 10px;
color: #555;
}}
.picture-coords {{
margin-top: 5px;
font-size: 0.8em;
color: #777;
}}
.no-pictures {{ .no-pictures {{
background-color: white; background-color: white;
padding: 20px; padding: 20px;
border-radius: 5px; border-radius: 5px;
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
text-align: center; text-align: center;
color: #777; color: #777;
}} }}
@ -300,42 +170,26 @@ def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir):
""" """
if pictures: if pictures:
html += """ <div class="gallery"> html += ' <div class="gallery">\n'
"""
for picture, file_path in zip(pictures, saved_files): for picture, file_path in zip(pictures, saved_files):
# Get relative path for the image
rel_path = file_path.name rel_path = file_path.name
html += f""" <div class="picture-card"> html += f""" <div class="picture-card">
<img src="{rel_path}" alt="Picture {picture['id']}"> <img src="{rel_path}" alt="Picture {picture['id']}">
<div class="picture-info"> <div class="picture-info">
<h3>Picture {picture['id']}</h3> <h3>Picture {picture['id']}</h3>
""" {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>
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>
</div> </div>
</div> </div>
""" """
html += """ </div> html += ' </div>\n'
"""
else: else:
html += """ <div class="no-pictures"> html += ' <div class="no-pictures">\n <h2>No pictures found on this page</h2>\n </div>\n'
<h2>No pictures found on this page</h2>
<p>The DocTags file doesn't contain any picture elements for this page.</p>
</div>
"""
html += """</body> html += '</body>\n</html>\n'
</html>
"""
# Save the HTML file
with open(index_file, 'w', encoding='utf-8') as f: with open(index_file, 'w', encoding='utf-8') as f:
f.write(html) f.write(html)
@ -343,10 +197,7 @@ def create_html_index(pictures, saved_files, pdf_name, page_num, output_dir):
return index_file return index_file
def main(): def main():
# Parse arguments
args = parse_arguments() args = parse_arguments()
# Create output directory
output_dir = ensure_results_folder(args.output) output_dir = ensure_results_folder(args.output)
try: try:
@ -361,48 +212,30 @@ def main():
print(f"Found {len(pictures)} picture elements.") print(f"Found {len(pictures)} picture elements.")
# Load the image from PDF # 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]}") 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: if args.adjust:
pictures = auto_adjust_coordinates(pictures, page_image.width, page_image.height) # Check if coordinates need normalization
elif args.scale != 1.0 or args.scale_x is not None or args.scale_y is not None: max_x = max([p['x2'] for p in pictures])
# Apply manual scaling max_y = max([p['y2'] for p in pictures])
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
print(f"Applying manual scaling: X={scale_x}, Y={scale_y}") if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
for picture in pictures: print(f"Detected normalized coordinates (0-{DEFAULT_GRID_SIZE} grid)")
picture['x1'] = int(picture['x1'] * scale_x) pictures = normalize_coordinates(pictures, page_image.width, page_image.height)
picture['y1'] = int(picture['y1'] * scale_y) else:
picture['x2'] = int(picture['x2'] * scale_x) pictures = auto_adjust_coordinates(pictures, page_image.width, page_image.height)
picture['y2'] = int(picture['y2'] * scale_y)
# Extract and save pictures # Extract and save pictures
saved_files = extract_and_save_pictures( saved_files = extract_and_save_pictures(
page_image, page_image, pictures, output_dir,
pictures, args.max_width, args.margin
output_dir,
args.max_width,
args.margin
) )
# Create HTML index # Create HTML index
pdf_name = Path(args.pdf).stem pdf_name = Path(args.pdf).stem
index_file = create_html_index(pictures, saved_files, pdf_name, args.page, output_dir) 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)}")
except Exception as e: except Exception as e:
print(f"Error: {e}") print(f"Error: {e}")

350
backend/page_treatment/visualizer.py
View file

@ -1,114 +1,44 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
""" """
DocTags Zone Visualizer - Simple script to visualize zones identified in DocTags format. DocTags Zone Visualizer - 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
""" """
import argparse import argparse
import os import os
import re import re
import sys
from pathlib import Path from pathlib import Path
from PIL import Image, ImageDraw 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 # Regular expression to extract location data
LOC_PATTERN = r'<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>' 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(): def parse_arguments():
"""Parse command line arguments.""" """Parse command line arguments."""
results_dir = ensure_results_folder() 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, parser.add_argument('--doctags', '-d', type=str, required=True,
help='Path to DocTags file') help='Path to DocTags file')
parser.add_argument('--pdf', '-p', type=str, required=True, parser.add_argument('--pdf', '-p', type=str, required=True,
help='Path to original PDF file') help='Path to original PDF file')
parser.add_argument('--page', type=int, default=8, parser.add_argument('--page', type=int, default=1,
help='Page number in PDF (starts at 1, default: 8)') help='Page number in PDF (starts at 1)')
parser.add_argument('--output', '-o', type=str, default=None, parser.add_argument('--output', '-o', type=str, default=None,
help='Output PNG file path (default: results/visualization_page_X.png)') help='Output PNG file path')
parser.add_argument('--dpi', type=int, default=200, parser.add_argument('--dpi', type=int, default=DEFAULT_DPI,
help='DPI for PDF rendering') 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', parser.add_argument('--adjust', action='store_true',
help='Try to automatically adjust scaling') help='Try to automatically adjust scaling')
return parser.parse_args() 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): def parse_doctags(doctags_path):
"""Parse DocTags file and extract zones with their coordinates.""" """Parse DocTags file and extract zones with their coordinates."""
if not os.path.exists(doctags_path): if not os.path.exists(doctags_path):
@ -118,86 +48,59 @@ def parse_doctags(doctags_path):
doctags_content = f.read() doctags_content = f.read()
# Extract content between <doctag> tags # Extract content between <doctag> tags
doctag_pattern = r'<doctag>(.*?)</doctag>' doctag_match = re.search(r'<doctag>(.*?)</doctag>', doctags_content, re.DOTALL)
doctag_match = re.search(doctag_pattern, doctags_content, re.DOTALL)
if not doctag_match: if not doctag_match:
raise ValueError("No <doctag> tags found in the file") raise ValueError("No <doctag> tags found in the file")
doctag_content = doctag_match.group(1) doctag_content = doctag_match.group(1)
# Find all tags with location information
zones = [] zones = []
# Find all tag starts # Find all tags with location information
tag_starts = re.finditer(r'<(\w+)>', doctag_content) tag_starts = re.finditer(r'<(\w+)>', doctag_content)
for tag_match in tag_starts: for tag_match in tag_starts:
tag_name = tag_match.group(1) tag_name = tag_match.group(1)
# Skip location tags themselves
if tag_name.startswith('loc_'): if tag_name.startswith('loc_'):
continue continue
# Find the end of the tag
tag_start_pos = tag_match.start() tag_start_pos = tag_match.start()
tag_end_pattern = f'</({tag_name})>' tag_end_pattern = f'</({tag_name})>'
tag_end_match = re.search(tag_end_pattern, doctag_content[tag_start_pos:]) tag_end_match = re.search(tag_end_pattern, doctag_content[tag_start_pos:])
if not tag_end_match: 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()] 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) loc_match = re.search(LOC_PATTERN, tag_content)
if loc_match: if loc_match:
# Extract coordinates
x1, y1, x2, y2 = map(int, loc_match.groups()) x1, y1, x2, y2 = map(int, loc_match.groups())
# Extract text content if available # Extract text content
text_content = ""
# Look for content between the location info and the closing tag
content_pattern = f'{LOC_PATTERN}(.*?)</{tag_name}>' content_pattern = f'{LOC_PATTERN}(.*?)</{tag_name}>'
content_match = re.search(content_pattern, tag_content, re.DOTALL) content_match = re.search(content_pattern, tag_content, re.DOTALL)
text_content = content_match.group(5).strip() if content_match else ""
if content_match:
text_content = content_match.group(5).strip()
zones.append({ zones.append({
'type': tag_name, 'type': tag_name,
'x1': x1, 'x1': x1, 'y1': y1,
'y1': y1, 'x2': x2, 'y2': y2,
'x2': x2,
'y2': y2,
'content': text_content 'content': text_content
}) })
return zones return zones
def create_debug_image(image, zones, page_num, output_path): def create_visualization(image, zones, page_num, output_path):
"""Create a debug image with rectangles around zones.""" """Create a visualization image with rectangles around zones."""
# Create a copy of the input image
debug_img = image.copy() debug_img = image.copy()
draw = ImageDraw.Draw(debug_img) 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 # Draw rectangles for each zone
for zone in zones: for zone in zones:
zone_type = zone['type'] 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( draw.rectangle(
[(zone['x1'], zone['y1']), (zone['x2'], zone['y2'])], [(zone['x1'], zone['y1']), (zone['x2'], zone['y2'])],
outline=color, outline=color,
@ -206,7 +109,7 @@ def create_debug_image(image, zones, page_num, output_path):
# Add zone type label # Add zone type label
label_width = len(zone_type) * 7 + 6 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( draw.rectangle(
[(label_x, zone['y1']), (label_x + label_width, zone['y1'] + 20)], [(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 fill=color
) )
# Draw page number on the debug image # Draw page number
draw.rectangle( draw.rectangle(
[(10, 10), (100, 40)], [(10, 10), (100, 40)],
fill=(0, 0, 0, 180), fill=(0, 0, 0, 180),
@ -231,167 +134,49 @@ def create_debug_image(image, zones, page_num, output_path):
fill=(255, 255, 255) fill=(255, 255, 255)
) )
# Save the debug image # Save the image
debug_img.save(output_path) debug_img.save(output_path)
print(f"Debug image saved to: {output_path}") print(f"Visualization saved to: {output_path}")
print(f"Absolute path: {output_path.absolute()}")
return debug_img return debug_img
def normalize_coordinates(zones, image_width, image_height, grid_size=500): def process_page(pdf_path, page_num, doctags_path, output_path, dpi, adjust):
"""
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):
"""Process a single page of the PDF with visualization.""" """Process a single page of the PDF with visualization."""
# Ensure results folder exists
results_dir = ensure_results_folder() results_dir = ensure_results_folder()
# Generate output path if not provided
if output_path is None: if output_path is None:
output_name = f"visualization_page_{page_num}.png" output_path = results_dir / f"visualization_page_{page_num}.png"
output_path = results_dir / output_name else:
output_path = Path(output_path)
# Make sure output_path is a Path object
output_path = Path(output_path)
# Load the page image # Load the page image
try: image = load_pdf_page(pdf_path, page_num, dpi)
image = load_image_from_pdf(pdf_path, page_num, dpi) print(f"Page {page_num} loaded: {image.size}")
print(f"Page {page_num} loaded: {image.size}")
except Exception as e:
print(f"Error loading page {page_num}: {e}")
return False
# Parse DocTags # Parse DocTags
try: zones = parse_doctags(doctags_path)
zones = parse_doctags(doctags_path) print(f"Found {len(zones)} zones in DocTags")
print(f"Found {len(zones)} zones in DocTags")
# Debug output to understand scaling issues if zones:
# After parsing the zones from DocTags # Check if we need to adjust coordinates
if zones: max_x = max([zone['x2'] for zone in zones])
img_width, img_height = image.size max_y = max([zone['y2'] for zone in zones])
print(f"Image dimensions: {img_width}x{img_height}")
# Check if we need to normalize grid coordinates # Auto-adjust if needed
max_x = max([zone['x2'] for zone in zones]) if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
max_y = max([zone['y2'] for zone in zones]) 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) # Create visualization
if max_x <= 500 and max_y <= 500: create_visualization(image, zones, page_num, output_path)
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)
return True return True
def main(): def main():
# Parse arguments
args = 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 # Check if files exist
if not os.path.exists(args.pdf): if not os.path.exists(args.pdf):
print(f"Error: PDF file not found: {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}") print(f"Error: DocTags file not found: {args.doctags}")
return return
# Process page(s) # Process the page
if args.page == 0: # Special case: process all pages process_page(
process_all_pages( args.pdf,
args.pdf, args.page,
args.doctags, args.doctags,
args.output, args.output,
args.dpi, args.dpi,
args.scale, args.adjust
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
)
if __name__ == "__main__": if __name__ == "__main__":
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 // DocTags Application State Management
const activeTasks = {}; const appState = {
let pollingInterval = null; activeTasks: {},
pollingInterval: null,
// Keep track of generated outputs generatedOutputs: {
const generatedOutputs = { visualizer: null,
visualizer: null, extractor: false
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() { function loadPDFPreview() {
const pdfFile = document.getElementById('pdf_file').value; const pdfFile = ui.getValue('pdf_file');
const pageNum = document.getElementById('page_num').value; const pageNum = ui.getValue('page_num');
if (!pdfFile) { if (!pdfFile) {
document.getElementById('pdf-preview-container').classList.add('hidden'); ui.hide('pdf-preview-container');
return; return;
} }
// Update page number display ui.setText('preview-page-num', pageNum);
document.getElementById('preview-page-num').textContent = 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'); const previewImg = document.getElementById('pdf-preview-image');
previewImg.src = `/pdf-preview/${encodeURIComponent(pdfFile)}/${pageNum}`; 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() { previewImg.onerror = function() {
this.alt = 'Failed to load PDF preview'; this.alt = 'Failed to load PDF preview';
console.error('Failed to load PDF preview'); console.error('Failed to load PDF preview');
}; };
} }
// Navigate to a specific page in the preview function changePreviewPage(delta) {
function navigateToPage(pageNum) { const currentPage = parseInt(ui.getValue('page_num'));
const pdfFile = document.getElementById('pdf_file').value; const newPage = Math.max(1, currentPage + delta);
if (!pdfFile) return; ui.setValue('page_num', newPage);
// Update the main page number input
document.getElementById('page_num').value = pageNum;
// Reload the preview
loadPDFPreview(); 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() { function goToPreviewPage() {
const pageInput = document.getElementById('preview-page-input'); const pageInput = document.getElementById('preview-page-input');
const pageNum = parseInt(pageInput.value); const pageNum = parseInt(pageInput.value);
if (pageNum && pageNum > 0) { if (pageNum && pageNum > 0) {
navigateToPage(pageNum); ui.setValue('page_num', pageNum);
loadPDFPreview();
} else { } else {
// Reset to current page if invalid pageInput.value = ui.getValue('page_num');
pageInput.value = document.getElementById('page_num').value;
} }
} }
// Tab switching functionality // Tab Management
function switchTab(tabName) { function switchTab(tabName) {
// Hide all tab contents // Hide all tab contents and deactivate buttons
const tabContents = document.querySelectorAll('.tab-content'); document.querySelectorAll('.tab-content').forEach(content => content.classList.remove('active'));
tabContents.forEach(content => content.classList.remove('active')); document.querySelectorAll('.tab-button').forEach(button => button.classList.remove('active'));
// Remove active class from all tab buttons // Show selected tab
const tabButtons = document.querySelectorAll('.tab-button');
tabButtons.forEach(button => button.classList.remove('active'));
// Show selected tab content
document.getElementById(tabName + '-tab').classList.add('active'); document.getElementById(tabName + '-tab').classList.add('active');
// Add active class to selected tab button
event.target.classList.add('active'); event.target.classList.add('active');
// Show previously generated content when switching tabs // Show previously generated content
if (tabName === 'analyzer') { if (tabName === 'analyzer') {
// Load PDF preview if a PDF is selected
loadPDFPreview(); loadPDFPreview();
} else if (tabName === 'visualizer' && generatedOutputs.visualizer) { } else if (tabName === 'visualizer' && appState.generatedOutputs.visualizer) {
document.getElementById('result-image').src = generatedOutputs.visualizer + '?t=' + new Date().getTime(); document.getElementById('result-image').src = appState.generatedOutputs.visualizer + '?t=' + Date.now();
document.getElementById('image-container').classList.remove('hidden'); ui.show('image-container');
} else if (tabName === 'extractor' && generatedOutputs.extractor) { } else if (tabName === 'extractor' && appState.generatedOutputs.extractor) {
loadExtractedImages(); loadExtractedImages();
} }
} }
// Update progress indicator // Progress Management
function updateProgress(completedSteps) { function updateProgress(completedSteps) {
for (let i = 1; i <= 3; i++) { for (let i = 1; i <= 3; i++) {
const stepIndicator = document.getElementById(`step-${i}`); const stepIndicator = document.getElementById(`step-${i}`);
@ -121,72 +141,202 @@ function updateProgress(completedSteps) {
} }
} }
// Load extracted images from the results folder // Task Management
function loadExtractedImages() { function startPolling() {
const imageGallery = document.getElementById('extracted-images-gallery'); if (!appState.pollingInterval) {
const imageContainer = document.getElementById('extracted-images-container'); appState.pollingInterval = setInterval(pollTasks, 1000);
}
// 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;
});
} }
// 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) { function openImageModal(imageSrc, caption, coords) {
// Create modal if it doesn't exist
let modal = document.getElementById('image-modal'); let modal = document.getElementById('image-modal');
if (!modal) { if (!modal) {
modal = document.createElement('div'); modal = document.createElement('div');
@ -205,403 +355,123 @@ function openImageModal(imageSrc, caption, coords) {
document.body.appendChild(modal); document.body.appendChild(modal);
} }
// Set image and info
document.getElementById('modal-image').src = imageSrc; document.getElementById('modal-image').src = imageSrc;
document.getElementById('modal-caption').textContent = caption; ui.setText('modal-caption', caption);
document.getElementById('modal-coords').textContent = coords; ui.setText('modal-coords', coords);
// Show modal
modal.classList.add('active'); modal.classList.add('active');
} }
// Close image modal
function closeImageModal() { function closeImageModal() {
const modal = document.getElementById('image-modal'); const modal = document.getElementById('image-modal');
if (modal) { if (modal) modal.classList.remove('active');
modal.classList.remove('active');
}
} }
// Load PDF files on page load // Environment Check
window.addEventListener('DOMContentLoaded', function() { async function checkEnvironment() {
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() {
const envDiv = document.getElementById('environment-check'); const envDiv = document.getElementById('environment-check');
const envDetails = document.getElementById('env-details'); const envDetails = document.getElementById('env-details');
envDiv.classList.remove('hidden'); ui.show('environment-check');
envDetails.innerHTML = '<div class="loader"></div> Checking environment...'; ui.setHtml('env-details', '<div class="loader"></div> Checking environment...');
fetch('/check-environment') try {
.then(response => response.json()) const data = await api.get('/check-environment');
.then(data => {
let html = '<ul>';
// Current working directory let html = '<ul>';
html += '<li>Working directory: <code>' + data.cwd + '</code></li>'; html += `<li>Working directory: <code>${data.cwd}</code></li>`;
html += `<li>Python version: <code>${data.python_version}</code></li>`;
// Python version if (data.missing_scripts.length === 0) {
html += '<li>Python version: <code>' + data.python_version + '</code></li>'; 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.pdf_files.length > 0) {
if (data.missing_scripts.length === 0) { html += `<li class="env-success">✓ Found ${data.pdf_files.length} PDF files</li>`;
html += '<li class="env-success">✓ All required scripts found</li>'; } else {
} else { html += '<li class="env-error">✗ No PDF files found in the working directory</li>';
html += '<li class="env-error">✗ Missing scripts: <code>' + data.missing_scripts.join(', ') + '</code></li>'; }
}
// PDF files check html += '</ul>';
if (data.pdf_files.length > 0) { ui.setHtml('env-details', html);
html += '<li class="env-success">✓ Found ' + data.pdf_files.length + ' PDF files: <code>' + data.pdf_files.join(', ') + '</code></li>'; } catch (error) {
} else { ui.setHtml('env-details', `<div class="env-error">Error checking environment: ${error.message}</div>`);
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>';
});
} }
// Manual command execution for debugging // Manual Command Execution
function manuallyRunScript() { async function manuallyRunScript() {
const command = prompt("Enter command to run (e.g., 'python backend/page_treatment/analyzer.py --image document.pdf --page 1')"); const command = prompt("Enter command to run:");
if (!command) return; if (!command) return;
const outputDiv = document.getElementById('output'); const outputDiv = document.getElementById('output');
outputDiv.textContent = 'Running command: ' + command + '\nPlease wait...'; ui.setText('output', 'Running command: ' + command + '\nPlease wait...');
outputDiv.classList.remove('hidden'); ui.show('output');
// Create form data
const formData = new FormData(); const formData = new FormData();
formData.append('command', command); formData.append('command', command);
// Send the command directly to backend try {
fetch('/run-manual-command', { const data = await api.post('/run-manual-command', formData);
method: 'POST', ui.setText('output',
body: formData 'Command: ' + command + '\n\n' +
}) (data.success ? 'Success!\n\n' : 'Failed!\n\n') +
.then(response => response.json()) (data.output || '') +
.then(data => { (data.error ? '\n\nError: ' + data.error : '')
outputDiv.textContent = 'Command: ' + command + '\n\n' + );
(data.success ? 'Success!\n\n' : 'Failed!\n\n') + } catch (error) {
(data.output || '') + ui.setText('output', 'Error running command: ' + error.message);
(data.error ? '\n\nError: ' + data.error : ''); }
}) }
.catch(error => {
outputDiv.textContent = '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 ui.setHtml('pdf-load-status', `<span class="success">Loaded ${data.length} PDF files</span>`);
function startPolling() { } catch (error) {
if (pollingInterval) { ui.setHtml('pdf-load-status', '<span class="error">Error: ' + error.message + '</span>');
return; // Already polling
} }
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 const previewPageInput = document.getElementById('preview-page-input');
function checkAndStopPolling() { if (previewPageInput) {
if (Object.keys(activeTasks).length === 0 && pollingInterval) { previewPageInput.addEventListener('keypress', function(e) {
clearInterval(pollingInterval); if (e.key === 'Enter') goToPreviewPage();
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');
}); });
} }
// 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 {}
});