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.tags import Tag, PhotoTag
|
||||||
from app.models.heaps import Heap, HeapPhoto
|
from app.models.heaps import Heap, HeapPhoto
|
||||||
from app.models.embeddings import Embedding
|
from app.models.embeddings import Embedding
|
||||||
|
from app.models.ocr_text import OCRText
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
'Photo',
|
'Photo',
|
||||||
@@ -15,5 +16,6 @@ __all__ = [
|
|||||||
'PhotoTag',
|
'PhotoTag',
|
||||||
'Heap',
|
'Heap',
|
||||||
'HeapPhoto',
|
'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:
|
except Exception as e:
|
||||||
logger.warning("Semantic search failed (models may not be loaded): %s", e)
|
logger.warning("Semantic search failed (models may not be loaded): %s", e)
|
||||||
|
|
||||||
# ── FTS search (placeholder for PR5) ──────────────────────────────
|
# ── FTS search (photos.search_vector + ocr_text) ────────────────
|
||||||
# Will be: tsvector @@ plainto_tsquery(q), ranked by ts_rank.
|
if q:
|
||||||
# For now, skip.
|
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 ────────────────────────────────────────────────────
|
# ── RRF fusion ────────────────────────────────────────────────────
|
||||||
k = 60
|
k = 60
|
||||||
|
|||||||
@@ -92,8 +92,44 @@ def vision_fanout(photo_id: str):
|
|||||||
|
|
||||||
@shared_task(name='ocr_photo', queue='vision')
|
@shared_task(name='ocr_photo', queue='vision')
|
||||||
def ocr_photo(photo_id: str):
|
def ocr_photo(photo_id: str):
|
||||||
"""OCR a photo — implemented in PR5."""
|
"""Run OCR on a photo and store text regions."""
|
||||||
return {'status': 'not_implemented'}
|
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')
|
@shared_task(name='detect_objects', queue='vision')
|
||||||
|
|||||||
Reference in New Issue
Block a user