108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
import os
|
|
import sys
|
|
if hasattr(sys.stdout, 'reconfigure'):
|
|
sys.stdout.reconfigure(encoding='utf-8')
|
|
if hasattr(sys.stderr, 'reconfigure'):
|
|
sys.stderr.reconfigure(encoding='utf-8')
|
|
|
|
import shutil
|
|
import argparse
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
import extractor
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Run the News-Scan end-to-end extraction pipeline.")
|
|
parser.add_argument(
|
|
"--max-pages",
|
|
type=int,
|
|
help="Limit processing to the first N pages."
|
|
)
|
|
parser.add_argument(
|
|
"--pages",
|
|
type=str,
|
|
help="Comma-separated list of specific 1-based page numbers to process (e.g. 1,3,5)."
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
input_dir = Path("input")
|
|
processed_dir = Path("processed")
|
|
output_dir = Path("output")
|
|
|
|
# Ensure directories exist
|
|
input_dir.mkdir(parents=True, exist_ok=True)
|
|
processed_dir.mkdir(parents=True, exist_ok=True)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Check for API key early
|
|
if not os.environ.get("ANTHROPIC_API_KEY"):
|
|
print("WARNING: 'ANTHROPIC_API_KEY' environment variable is not set!")
|
|
print("The political filter and AI insights generation will be SKIPPED.")
|
|
print("To generate the final PDF report, you must set this key.")
|
|
print("-" * 60)
|
|
|
|
# Find the first PDF in the input folder
|
|
pdf_files = list(input_dir.glob("*.pdf"))
|
|
if not pdf_files:
|
|
print(f"Error: No PDF files found in {input_dir.absolute()}")
|
|
return
|
|
|
|
pdf_path = pdf_files[0]
|
|
print(f"Found file to process: {pdf_path.name}")
|
|
|
|
job_id = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
job_dir = output_dir / job_id
|
|
job_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Parse selected pages
|
|
pages_to_process = None
|
|
if args.pages:
|
|
try:
|
|
pages_to_process = [int(p.strip()) for p in args.pages.split(",") if p.strip()]
|
|
print(f"Limiting processing to specific pages: {pages_to_process}")
|
|
except ValueError:
|
|
print("Error: --pages must be a comma-separated list of integers (e.g., 1,3,5)")
|
|
return
|
|
elif args.max_pages:
|
|
print(f"Limiting processing to first {args.max_pages} pages")
|
|
else:
|
|
print("Processing all pages of the document")
|
|
|
|
print(f"Starting full pipeline extraction on {pdf_path}")
|
|
print(f"Output directory: {job_dir}")
|
|
print()
|
|
|
|
# Run the full pipeline (extract, group, crop, OCR, filter)
|
|
result = extractor.process_pdf(
|
|
str(pdf_path),
|
|
str(job_dir),
|
|
4200,
|
|
7400,
|
|
job_id=job_id,
|
|
max_pages=args.max_pages,
|
|
pages_to_process=pages_to_process
|
|
)
|
|
print(f"\nExtraction complete! Found {result['articles']} articles across {result['pages']} pages.")
|
|
|
|
# Move the file from input to processed
|
|
dest_path = processed_dir / pdf_path.name
|
|
|
|
# Ensure no collision in processed folder
|
|
if dest_path.exists():
|
|
timestamp = datetime.now().strftime("%H%M%S")
|
|
dest_path = processed_dir / f"{pdf_path.stem}_{timestamp}{pdf_path.suffix}"
|
|
|
|
shutil.move(str(pdf_path.absolute()), str(dest_path.absolute()))
|
|
print(f"\nMoved processed file to: {dest_path}")
|
|
|
|
# Check if the PDF report was generated
|
|
pdf_report = job_dir / "political_articles_insights.pdf"
|
|
if pdf_report.exists():
|
|
print(f"\nSUCCESS: AI Insights PDF generated at {pdf_report}")
|
|
else:
|
|
print(f"\nNOTE: The AI Insights PDF was not generated. This usually means no political articles were found, or the ANTHROPIC_API_KEY was missing.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|