feat(ml): image_prediction table + backfill + dual-write (#768 step 1)
Normalize tagger predictions out of the image_record.tagger_predictions JSON blob into a queryable per-prediction table. Step 1 of the cutover (expand): additive + low-risk — reads still use the JSON, this just adds the table and keeps it populated. - ImagePrediction(image_record_id, raw_name, category, score) — stores the RAW tagger vocab name (not tag_id) so read-time alias→canonical resolution is unchanged. Indexed for per-image reads + by (raw_name, score). - Migration 0045: create table + set-based backfill from the JSON via json_each (fast post-#764-prune). The old column stays (vestigial) and is dropped in a later follow-up — DROP needs an ACCESS EXCLUSIVE lock on the hot image_record table, so it waits for a quiesced-worker window. - tag_and_embed dual-writes the rows (delete-then-insert, idempotent); tagger_store_floor already applied in infer(). Next: switch suggestion + allowlist reads to the table, then drop the JSON write. Plan-task #768. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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")
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
+19
-2
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user