50 lines
2.4 KiB
Python
50 lines
2.4 KiB
Python
from app.pipeline.pdf_reader import PDFReader
|
|
from app.pipeline.layout_detector import LayoutDetector
|
|
from app.pipeline.article_segmenter import ArticleSegmenter
|
|
from app.pipeline.ocr_engine import OCREngine
|
|
from app.pipeline.headline import detect_headline
|
|
from app.pipeline.cropper import crop_article, crop_figures
|
|
from app.pipeline.categorizer import Categorizer
|
|
from app.pipeline.enricher import Enricher
|
|
from app.config import settings
|
|
import os
|
|
|
|
class Pipeline:
|
|
def __init__(self):
|
|
self.layout = LayoutDetector()
|
|
self.segmenter = ArticleSegmenter()
|
|
self.ocr = OCREngine(settings.OCR_LANGS)
|
|
self.cat = Categorizer()
|
|
self.enricher = Enricher()
|
|
|
|
def process(self, pdf_path, newspaper, edition, date, job_dir):
|
|
reader = PDFReader(pdf_path)
|
|
results = []
|
|
for page in reader.iter_pages():
|
|
h, w = page["image"].shape[:2]
|
|
blocks = self.layout.detect(page["image"])
|
|
articles = self.segmenter.segment(blocks, w, h)
|
|
for art in articles:
|
|
text, conf = self.ocr.extract_text(page, art["bbox"], page["native_words"])
|
|
if len(text.strip()) < 80: # skip noise/ads
|
|
continue
|
|
headline = detect_headline(art, page, self.ocr)
|
|
img_path = crop_article(page["image"], art["bbox"], job_dir)
|
|
fig_paths = crop_figures(page["image"], art, job_dir)
|
|
meta = self.cat.classify(headline, text)
|
|
embedding = self.enricher.embed(headline + " " + text[:1000])
|
|
results.append({
|
|
"date": date, "newspaper": newspaper, "edition": edition,
|
|
"page": page["page_number"], "headline": headline,
|
|
"category": meta["category"], "subcategory": meta.get("subcategory"),
|
|
"text": text, "summary": meta.get("summary"),
|
|
"keywords": meta.get("keywords"), "persons": meta.get("persons"),
|
|
"locations": meta.get("locations"), "organizations": meta.get("organizations"),
|
|
"sentiment": meta.get("sentiment"), "language": "en",
|
|
"image_path": img_path, "figure_paths": fig_paths,
|
|
"bounding_box": art["bbox"],
|
|
"confidence": round(art["confidence"] * meta.get("confidence", 1), 3),
|
|
"embedding": embedding,
|
|
})
|
|
return results
|