feat: add OCR text extraction and Postgres full-text search
- Create ocr_text table for storing per-region OCR results - Add tsvector search_vector column to photos with GIN index and auto-update trigger on filename/user_title/user_notes - Implement ocr_photo Celery task using rapidocr-onnxruntime - Add FTS leg to hybrid search: queries photos.search_vector and ocr_text via UNION, fused with semantic results via RRF (k=60) Migration 0004 backfills search_vector for existing rows. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
82
backend/alembic/versions/0004_ocr_text_and_fts.py
Normal file
82
backend/alembic/versions/0004_ocr_text_and_fts.py
Normal file
@@ -0,0 +1,82 @@
|
||||
"""ocr_text table and Postgres FTS
|
||||
|
||||
Revision ID: 0004_ocr_fts
|
||||
Revises: 0003_pgvector_embeddings
|
||||
Create Date: 2026-04-10
|
||||
|
||||
Create ocr_text table for storing OCR results. Add a tsvector column
|
||||
to photos for unified full-text search (filename + user_title +
|
||||
user_notes) with a GIN index. OCR text is rolled up into a materialized
|
||||
view or joined at query time.
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0004_ocr_fts"
|
||||
down_revision: Union[str, None] = "0003_pgvector_embeddings"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ── ocr_text table ────────────────────────────────────────────────
|
||||
op.execute("""
|
||||
CREATE TABLE IF NOT EXISTS ocr_text (
|
||||
id VARCHAR PRIMARY KEY,
|
||||
photo_id VARCHAR NOT NULL REFERENCES photos(id) ON DELETE CASCADE,
|
||||
text TEXT NOT NULL,
|
||||
language VARCHAR(8) DEFAULT '',
|
||||
confidence FLOAT,
|
||||
bbox JSONB,
|
||||
created_at TIMESTAMPTZ DEFAULT now()
|
||||
)
|
||||
""")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_ocr_text_photo_id ON ocr_text(photo_id)")
|
||||
|
||||
# ── tsvector column on photos ─────────────────────────────────────
|
||||
op.execute("ALTER TABLE photos ADD COLUMN IF NOT EXISTS search_vector tsvector")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS ix_photos_search_vector ON photos USING GIN (search_vector)")
|
||||
|
||||
# Trigger to auto-update search_vector on INSERT/UPDATE
|
||||
op.execute("""
|
||||
CREATE OR REPLACE FUNCTION photos_search_vector_update() RETURNS trigger AS $$
|
||||
BEGIN
|
||||
NEW.search_vector :=
|
||||
setweight(to_tsvector('english', coalesce(NEW.filename, '')), 'A') ||
|
||||
setweight(to_tsvector('english', coalesce(NEW.user_title, '')), 'A') ||
|
||||
setweight(to_tsvector('english', coalesce(NEW.user_notes, '')), 'B');
|
||||
RETURN NEW;
|
||||
END
|
||||
$$ LANGUAGE plpgsql;
|
||||
""")
|
||||
op.execute("""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_trigger WHERE tgname = 'photos_search_vector_trigger'
|
||||
) THEN
|
||||
CREATE TRIGGER photos_search_vector_trigger
|
||||
BEFORE INSERT OR UPDATE OF filename, user_title, user_notes
|
||||
ON photos
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION photos_search_vector_update();
|
||||
END IF;
|
||||
END $$;
|
||||
""")
|
||||
|
||||
# Backfill existing rows
|
||||
op.execute("""
|
||||
UPDATE photos SET search_vector =
|
||||
setweight(to_tsvector('english', coalesce(filename, '')), 'A') ||
|
||||
setweight(to_tsvector('english', coalesce(user_title, '')), 'A') ||
|
||||
setweight(to_tsvector('english', coalesce(user_notes, '')), 'B')
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TRIGGER IF EXISTS photos_search_vector_trigger ON photos")
|
||||
op.execute("DROP FUNCTION IF EXISTS photos_search_vector_update()")
|
||||
op.execute("DROP INDEX IF EXISTS ix_photos_search_vector")
|
||||
op.execute("ALTER TABLE photos DROP COLUMN IF EXISTS search_vector")
|
||||
op.execute("DROP TABLE IF EXISTS ocr_text")
|
||||
@@ -6,6 +6,7 @@ from app.models.folders import Folder, SourceRoot
|
||||
from app.models.tags import Tag, PhotoTag
|
||||
from app.models.heaps import Heap, HeapPhoto
|
||||
from app.models.embeddings import Embedding
|
||||
from app.models.ocr_text import OCRText
|
||||
|
||||
__all__ = [
|
||||
'Photo',
|
||||
@@ -15,5 +16,6 @@ __all__ = [
|
||||
'PhotoTag',
|
||||
'Heap',
|
||||
'HeapPhoto',
|
||||
'Embedding'
|
||||
'Embedding',
|
||||
'OCRText',
|
||||
]
|
||||
20
backend/app/models/ocr_text.py
Normal file
20
backend/app/models/ocr_text.py
Normal file
@@ -0,0 +1,20 @@
|
||||
"""
|
||||
OCR text model — stores text regions extracted from photos via rapidocr.
|
||||
"""
|
||||
from sqlalchemy import Column, String, Float, ForeignKey, Text, DateTime, func
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
import uuid
|
||||
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class OCRText(Base):
|
||||
__tablename__ = 'ocr_text'
|
||||
|
||||
id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
|
||||
photo_id = Column(String, ForeignKey('photos.id', ondelete='CASCADE'), nullable=False, index=True)
|
||||
text = Column(Text, nullable=False)
|
||||
language = Column(String(8), default='')
|
||||
confidence = Column(Float)
|
||||
bbox = Column(JSONB) # [x1, y1, x2, y2] normalized 0-1
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
@@ -62,9 +62,30 @@ async def hybrid_search(
|
||||
except Exception as e:
|
||||
logger.warning("Semantic search failed (models may not be loaded): %s", e)
|
||||
|
||||
# ── FTS search (placeholder for PR5) ──────────────────────────────
|
||||
# Will be: tsvector @@ plainto_tsquery(q), ranked by ts_rank.
|
||||
# For now, skip.
|
||||
# ── FTS search (photos.search_vector + ocr_text) ────────────────
|
||||
if q:
|
||||
try:
|
||||
fts_stmt = text("""
|
||||
SELECT id, ts_rank(search_vector, plainto_tsquery('english', :q)) AS rank
|
||||
FROM photos
|
||||
WHERE search_vector @@ plainto_tsquery('english', :q)
|
||||
UNION
|
||||
SELECT o.photo_id AS id,
|
||||
MAX(o.confidence) AS rank
|
||||
FROM ocr_text o
|
||||
WHERE to_tsvector('english', o.text) @@ plainto_tsquery('english', :q)
|
||||
GROUP BY o.photo_id
|
||||
ORDER BY rank DESC
|
||||
LIMIT 200
|
||||
""")
|
||||
fts_rows = (await db.execute(fts_stmt, {"q": q})).fetchall()
|
||||
for rank, (photo_id, score) in enumerate(fts_rows):
|
||||
if photo_id not in results:
|
||||
results[photo_id] = {"semantic_rank": None, "fts_rank": rank}
|
||||
else:
|
||||
results[photo_id]["fts_rank"] = rank
|
||||
except Exception as e:
|
||||
logger.warning("FTS search failed: %s", e)
|
||||
|
||||
# ── RRF fusion ────────────────────────────────────────────────────
|
||||
k = 60
|
||||
|
||||
@@ -92,8 +92,44 @@ def vision_fanout(photo_id: str):
|
||||
|
||||
@shared_task(name='ocr_photo', queue='vision')
|
||||
def ocr_photo(photo_id: str):
|
||||
"""OCR a photo — implemented in PR5."""
|
||||
return {'status': 'not_implemented'}
|
||||
"""Run OCR on a photo and store text regions."""
|
||||
if not settings.vision.enabled or not settings.vision.ocr.enabled:
|
||||
return {'status': 'skipped', 'reason': 'OCR disabled'}
|
||||
return asyncio.run(_ocr_photo_async(photo_id))
|
||||
|
||||
|
||||
async def _ocr_photo_async(photo_id: str):
|
||||
image = _load_thumb(photo_id, "large") # 1280px for better OCR accuracy
|
||||
if image is None:
|
||||
return {'status': 'error', 'message': 'thumbnail not found'}
|
||||
|
||||
from app.services.vision.registry import registry
|
||||
ocr_engine = registry.get_ocr()
|
||||
results = ocr_engine.run(image)
|
||||
|
||||
if not results:
|
||||
logger.info("No OCR text found for photo %s", photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'regions': 0}
|
||||
|
||||
from app.models.ocr_text import OCRText
|
||||
|
||||
async with AsyncSessionLocal() as session:
|
||||
# Delete existing OCR results for this photo (re-run safe)
|
||||
await session.execute(
|
||||
delete(OCRText).where(OCRText.photo_id == photo_id)
|
||||
)
|
||||
for r in results:
|
||||
session.add(OCRText(
|
||||
photo_id=photo_id,
|
||||
text=r.text,
|
||||
language=r.language,
|
||||
confidence=r.confidence,
|
||||
bbox=r.bbox,
|
||||
))
|
||||
await session.commit()
|
||||
|
||||
logger.info("OCR: %d text regions for photo %s", len(results), photo_id)
|
||||
return {'status': 'success', 'photo_id': photo_id, 'regions': len(results)}
|
||||
|
||||
|
||||
@shared_task(name='detect_objects', queue='vision')
|
||||
|
||||
Reference in New Issue
Block a user