89 lines
4.3 KiB
Python
89 lines
4.3 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)
|
|
|
|
# Semantic Embedding for RAG
|
|
raw_text_for_embed = headline + " " + text[:1000]
|
|
embedding = self.enricher.embed(raw_text_for_embed)
|
|
|
|
# Fetch few-shot examples from DB
|
|
few_shot_examples = []
|
|
try:
|
|
from app.db.session import SessionLocal
|
|
from sqlalchemy import text as sqla_text
|
|
db = SessionLocal()
|
|
# Query for up to 3 verified examples using cosine distance
|
|
rows = db.execute(sqla_text("""
|
|
SELECT headline, text, category, subcategory, summary, sentiment, keywords, persons, locations, organizations
|
|
FROM articles
|
|
WHERE human_verified = TRUE
|
|
ORDER BY embedding <=> :v LIMIT 3
|
|
"""), {"v": str(embedding)}).fetchall()
|
|
for r in rows:
|
|
ex = {
|
|
"headline": r.headline,
|
|
"text": r.text,
|
|
"classification": {
|
|
"category": r.category,
|
|
"subcategory": r.subcategory,
|
|
"summary": r.summary,
|
|
"sentiment": r.sentiment,
|
|
"keywords": r.keywords,
|
|
"persons": r.persons,
|
|
"locations": r.locations,
|
|
"organizations": r.organizations
|
|
}
|
|
}
|
|
few_shot_examples.append(ex)
|
|
except Exception as e:
|
|
print(f"RAG fetch failed: {e}")
|
|
finally:
|
|
if 'db' in locals(): db.close()
|
|
|
|
meta = self.cat.classify(headline, text, few_shot_examples=few_shot_examples)
|
|
|
|
results.append({
|
|
"date": date, "newspaper": newspaper, "edition": edition,
|
|
"page": page["page_number"], "headline": headline,
|
|
"category": meta.get("category", "Others"), "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
|