diff --git a/alembic/versions/0045_image_prediction_table.py b/alembic/versions/0045_image_prediction_table.py new file mode 100644 index 0000000..b3bec55 --- /dev/null +++ b/alembic/versions/0045_image_prediction_table.py @@ -0,0 +1,73 @@ +"""image_prediction table + backfill from image_record.tagger_predictions + +Normalizes the per-image tagger predictions out of the JSON blob into a +queryable table (#768). Backfills from the existing JSON in one set-based +INSERT…SELECT over json_each — fast because the #764 prune already shrank +each row to its >=0.70 entries. The old image_record.tagger_predictions +column is left in place here (vestigial) and dropped in a follow-up once the +code cutover is verified — dropping it needs an ACCESS EXCLUSIVE lock on the +hot image_record table (the 0044 lock class), so it's deferred to a +quiesced-worker window. + +Revision ID: 0045 +Revises: 0044 +Create Date: 2026-06-10 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0045" +down_revision: Union[str, None] = "0044" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "image_prediction", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "image_record_id", sa.Integer(), + sa.ForeignKey("image_record.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("raw_name", sa.String(length=255), nullable=False), + sa.Column("category", sa.String(length=64), nullable=False), + sa.Column("score", sa.Float(), nullable=False), + sa.UniqueConstraint( + "image_record_id", "raw_name", name="image_raw_name", + ), + ) + op.create_index( + "ix_image_prediction_image", "image_prediction", ["image_record_id"], + ) + op.create_index( + "ix_image_prediction_name_score", "image_prediction", + ["raw_name", "score"], + ) + # Backfill from the JSON blob. json_each expands {name: {category, + # confidence}} into one row per prediction. category defaults to 'general' + # to mirror the suggestion read path; rows with no confidence are skipped. + op.execute( + """ + INSERT INTO image_prediction (image_record_id, raw_name, category, score) + SELECT ir.id, + je.key, + COALESCE(je.value ->> 'category', 'general'), + (je.value ->> 'confidence')::double precision + FROM image_record ir, + json_each(ir.tagger_predictions) je + WHERE ir.tagger_predictions IS NOT NULL + AND je.value ->> 'confidence' IS NOT NULL + ON CONFLICT (image_record_id, raw_name) DO NOTHING + """ + ) + + +def downgrade() -> None: + op.drop_index("ix_image_prediction_name_score", "image_prediction") + op.drop_index("ix_image_prediction_image", "image_prediction") + op.drop_table("image_prediction") diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 00047f9..19d444b 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -7,6 +7,7 @@ from .backup_run import BackupRun from .base import Base from .credential import Credential from .download_event import DownloadEvent +from .image_prediction import ImagePrediction from .image_provenance import ImageProvenance from .image_record import ImageRecord from .import_batch import ImportBatch @@ -45,6 +46,7 @@ __all__ = [ "SeriesPage", "SeriesSuggestion", "ImageRecord", + "ImagePrediction", "ImageProvenance", "Tag", "TagKind", diff --git a/backend/app/models/image_prediction.py b/backend/app/models/image_prediction.py new file mode 100644 index 0000000..f532243 --- /dev/null +++ b/backend/app/models/image_prediction.py @@ -0,0 +1,37 @@ +"""ImagePrediction — one row per (image, tagger vocab prediction). + +Replaces the image_record.tagger_predictions JSON blob (#768). Storing the +raw Camie/booru vocab name (not a tag_id) preserves the suggestion read +path's semantics: raw_name → canonical Tag resolution happens at read time +via the alias map, and accepting a prediction can CREATE the Tag. The store +floor (ml_settings.tagger_store_floor) is applied at WRITE time, so only +predictions >= the floor land here. +""" + +from sqlalchemy import Float, ForeignKey, Index, String, UniqueConstraint +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class ImagePrediction(Base): + __tablename__ = "image_prediction" + __table_args__ = ( + UniqueConstraint( + "image_record_id", "raw_name", name="image_raw_name", + ), + # Per-image read (suggestion build) and the "images with tag X above + # Y" query the JSON blob never allowed. + Index("ix_image_prediction_image", "image_record_id"), + Index("ix_image_prediction_name_score", "raw_name", "score"), + ) + + id: Mapped[int] = mapped_column(primary_key=True) + image_record_id: Mapped[int] = mapped_column( + ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, + ) + # The raw tagger vocab key (booru form) — NOT a tag_id. Resolved to a + # canonical Tag at read time, exactly as the old JSON keys were. + raw_name: Mapped[str] = mapped_column(String(255), nullable=False) + category: Mapped[str] = mapped_column(String(64), nullable=False) + score: Mapped[float] = mapped_column(Float, nullable=False) diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index ff09258..c79ed8d 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -10,11 +10,11 @@ import logging from pathlib import Path from celery.exceptions import SoftTimeLimitExceeded -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.exc import DBAPIError, OperationalError from ..celery_app import celery -from ..models import ImageRecord, MLSettings +from ..models import ImagePrediction, ImageRecord, MLSettings from ._sync_engine import sync_session_factory as _sync_session_factory log = logging.getLogger(__name__) @@ -162,6 +162,23 @@ def tag_and_embed(self, image_id: int) -> dict: record.siglip_embedding = embedding.tolist() record.siglip_model_version = settings.embedder_model_version session.add(record) + # Write the normalized image_prediction rows (#768). Delete-then- + # insert keeps a re-tag idempotent. tagger_store_floor was already + # applied in tagger.infer, so preds is the >=floor set. (Transitional + # dual-write alongside the JSON column until the read cutover lands.) + session.execute( + delete(ImagePrediction).where( + ImagePrediction.image_record_id == image_id + ) + ) + session.add_all([ + ImagePrediction( + image_record_id=image_id, raw_name=name, + category=p.get("category", "general"), + score=float(p.get("confidence", 0.0)), + ) + for name, p in preds.items() + ]) session.commit() except SoftTimeLimitExceeded: log.error( diff --git a/tests/test_image_prediction.py b/tests/test_image_prediction.py new file mode 100644 index 0000000..43e05e9 --- /dev/null +++ b/tests/test_image_prediction.py @@ -0,0 +1,57 @@ +"""#768: image_prediction table — model + constraints round-trip.""" +import pytest +from sqlalchemy import select +from sqlalchemy.exc import IntegrityError + +from backend.app.models import ImagePrediction, ImageRecord + +pytestmark = pytest.mark.integration + + +async def _make_image(db, path="/img/p0.jpg", sha="0"): + rec = ImageRecord( + path=path, sha256=sha.ljust(64, "0")[:64], size_bytes=10, + mime="image/jpeg", origin="imported_filesystem", + ) + db.add(rec) + await db.flush() + return rec + + +@pytest.mark.asyncio +async def test_image_prediction_round_trip(db): + rec = await _make_image(db) + db.add_all([ + ImagePrediction( + image_record_id=rec.id, raw_name="blue_eyes", + category="general", score=0.92, + ), + ImagePrediction( + image_record_id=rec.id, raw_name="hatsune_miku", + category="character", score=0.88, + ), + ]) + await db.commit() + + rows = (await db.execute( + select(ImagePrediction.raw_name, ImagePrediction.score) + .where(ImagePrediction.image_record_id == rec.id) + .order_by(ImagePrediction.score.desc()) + )).all() + assert [r.raw_name for r in rows] == ["blue_eyes", "hatsune_miku"] + + +@pytest.mark.asyncio +async def test_image_prediction_unique_per_image_name(db): + rec = await _make_image(db, path="/img/p1.jpg", sha="1") + db.add(ImagePrediction( + image_record_id=rec.id, raw_name="dup", + category="general", score=0.9, + )) + await db.commit() + db.add(ImagePrediction( + image_record_id=rec.id, raw_name="dup", + category="general", score=0.7, + )) + with pytest.raises(IntegrityError): + await db.commit()