diff --git a/backend/apps/library/tests/test_uploads.py b/backend/apps/library/tests/test_uploads.py index b38dd27..e8a6f4d 100644 --- a/backend/apps/library/tests/test_uploads.py +++ b/backend/apps/library/tests/test_uploads.py @@ -1,7 +1,8 @@ -"""Staged uploads: complete-set listing, bulk rating and IQDB recording.""" +"""Staged uploads: listing, phases (MD5/visual/IQDB) and the bulk tool.""" import base64 import hashlib +import io import json import shutil import tempfile @@ -11,6 +12,8 @@ from django.contrib.auth import get_user_model from django.core.files.uploadedfile import SimpleUploadedFile from django.test import Client, TestCase, override_settings +from PIL import Image + from rest_framework.authtoken.models import Token from apps.library.models import MediaItem, TempUpload @@ -23,6 +26,13 @@ TINY_PNG = base64.b64decode( ) +def similar_png(color=(200, 30, 40)): + """A 1x1 PNG with the same visual hashes but a different MD5.""" + buffer = io.BytesIO() + Image.new("RGB", (1, 1), color).save(buffer, format="PNG") + return buffer.getvalue() + + def jpost(client, path, body=None): return client.post(path, data=json.dumps(body or {}), content_type="application/json") @@ -85,7 +95,7 @@ class TempUploadListTests(TestCase): self.assertEqual(Client().get("/api/uploads/").status_code, 401) -class BulkResolveTests(TestCase): +class StagedUploadWorkflowTests(TestCase): @classmethod def setUpClass(cls): super().setUpClass() @@ -196,6 +206,61 @@ class BulkResolveTests(TestCase): self.assertEqual(body["errors"][0]["error"], "already in the library") self.assertEqual(MediaItem.objects.count(), 1) + def upload_via_api(self, client, label="phase.png"): + return client.post( + "/api/uploads/", + { + "file": SimpleUploadedFile( + label, similar_png(), content_type="image/png" + ) + }, + ) + + def seed_library_item(self, client, label): + seed = self.make_temp(self.uploader, label) + response = self.resolve_bulk(client, [seed.id], "s") + self.assertEqual(response.status_code, 200) + return seed + + def test_upload_defers_visual_similarity_to_its_phase(self): + client = self.api_client(self.uploader) + self.seed_library_item(client, "seed-defer") + response = self.upload_via_api(client) + self.assertEqual(response.status_code, 201) + body = response.json() + # Uploading must not run the expensive library scan. + self.assertEqual(body["status"], TempUpload.STATUS_PENDING) + self.assertFalse(body["visual_matches"]) + + def test_visual_match_phase_flags_similar_library_items(self): + client = self.api_client(self.uploader) + self.seed_library_item(client, "seed-visual") + temp = self.make_temp(self.uploader, "check-visual") + response = client.post(f"/api/uploads/{temp.id}/visual-match/") + self.assertEqual(response.status_code, 200) + body = response.json() + self.assertEqual(body["status"], TempUpload.STATUS_VISUAL_MATCH) + self.assertGreaterEqual(len(body["visual_matches"]), 1) + temp.refresh_from_db() + self.assertEqual(temp.status, TempUpload.STATUS_VISUAL_MATCH) + + def test_visual_match_phase_with_empty_library_stays_pending(self): + client = self.api_client(self.uploader) + temp = self.make_temp(self.uploader, "no-matches") + response = client.post(f"/api/uploads/{temp.id}/visual-match/") + self.assertEqual(response.status_code, 200) + body = response.json() + self.assertEqual(body["visual_matches"], []) + self.assertEqual(body["status"], TempUpload.STATUS_PENDING) + + def test_visual_match_phase_rejects_completed_uploads(self): + client = self.api_client(self.uploader) + temp = self.make_temp( + self.uploader, "done-visual", status=TempUpload.STATUS_COMPLETED + ) + response = client.post(f"/api/uploads/{temp.id}/visual-match/") + self.assertEqual(response.status_code, 400) + class IqdbRecordingTests(TestCase): """The modal needs to tell "checked, no match" from "never checked".""" diff --git a/backend/apps/library/uploads.py b/backend/apps/library/uploads.py index 7ec5a03..6bae436 100644 --- a/backend/apps/library/uploads.py +++ b/backend/apps/library/uploads.py @@ -190,18 +190,37 @@ class TempUploadViewSet( temp.resolution = TempUpload.RESOLUTION_DUPLICATE temp.library_item = existing temp.file.delete(save=False) - else: - matches = find_library_matches( - temp.file.path, user=request.user, request=request - ) - if matches: - temp.visual_matches = matches - temp.status = TempUpload.STATUS_VISUAL_MATCH + # Visual similarity is deliberately a separate phase (the + # /visual-match action) so a large batch uploads at full speed and + # the board runs MD5 -> visual -> IQDB over the whole batch. temp.save() return Response( self.get_serializer(temp).data, status=status.HTTP_201_CREATED ) + @action(detail=True, methods=["post"], url_path="visual-match") + def visual_match(self, request, pk=None): + """Run the local visual-similarity pass for one staged upload.""" + temp = self.get_object() + if temp.status == TempUpload.STATUS_COMPLETED: + return Response( + {"detail": "This upload is already in the library."}, + status=status.HTTP_400_BAD_REQUEST, + ) + if not temp.file: + return Response( + {"detail": "The staged file is missing."}, + status=status.HTTP_400_BAD_REQUEST, + ) + matches = find_library_matches( + temp.file.path, user=request.user, request=request + ) + temp.visual_matches = matches + if matches and temp.status == TempUpload.STATUS_PENDING: + temp.status = TempUpload.STATUS_VISUAL_MATCH + temp.save(update_fields=["visual_matches", "status", "updated_at"]) + return Response(self.get_serializer(temp).data) + @action(detail=True, methods=["get", "head"], permission_classes=[AllowAny]) def file(self, request, pk=None): """Serve the staged file; accepts a signed URL for media tags.""" diff --git a/frontend/src/features/upload/UploadPage.tsx b/frontend/src/features/upload/UploadPage.tsx index 86ee394..9ad34b6 100644 --- a/frontend/src/features/upload/UploadPage.tsx +++ b/frontend/src/features/upload/UploadPage.tsx @@ -127,6 +127,32 @@ function isFatalIqdbError(error: unknown): boolean { return error instanceof TypeError || error instanceof DOMException; } +/** Small labeled progress bar for one upload-pipeline phase. */ +function PhaseBar({ + label, + percent, + tone, +}: { + label: string; + percent: number; + tone: string; +}) { + return ( +
+

+ + {label} +

+
+
+
+
+ ); +} + function TempCard({ temp, checking, @@ -708,6 +734,11 @@ export default function UploadPage() { const [dragging, setDragging] = useState(false); const [entries, setEntries] = useState([]); const [batch, setBatch] = useState({ total: 0, done: 0, failed: 0 }); + const [pipeline, setPipeline] = useState<{ + phase: "md5" | "visual"; + done: number; + total: number; + } | null>(null); const [busy, setBusy] = useState(false); const [processing, setProcessing] = useState(false); const [checkingIds, setCheckingIds] = useState([]); @@ -931,48 +962,83 @@ export default function UploadPage() { } } - async function autoMatchAndCheck(created: TempUpload[]) { + /** + * Phased processing for one finished upload batch: every file goes through + * the e621 MD5 lookup (batched), then local visual similarity, then IQDB. + * Phases never interleave per file, and each step lands on the board as it + * finishes. + */ + async function runUploadPipeline(created: TempUpload[]) { const creds = effectiveCredentials(credentials); - const unresolved = created.filter((temp) => temp.status === "pending"); + const targets = created.filter((temp) => temp.status === "pending"); + if (targets.length === 0) { + invalidateUploads(); + return; + } + const targetIds = new Set(targets.map((temp) => temp.temp_id)); - // 1) Auto-complete files whose MD5 matches an e621 post. The original app - // sent 75 md5: metatags per posts.json query. - if (unresolved.length > 0) { - try { - for ( - let index = 0; - index < unresolved.length; - index += MD5_BATCH_SIZE - ) { - const md5Batch = unresolved.slice(index, index + MD5_BATCH_SIZE); - const posts = await fetchPostsByMd5( - creds, - md5Batch.map((temp) => temp.md5), - ); - for (const post of posts) { - const match = md5Batch.find((temp) => temp.md5 === post.file.md5); - if (!match) continue; - try { - const updated = await api( - `/api/uploads/${match.temp_id}/resolve/`, - { method: "POST", json: { mode: "link", post, auto: true } }, - ); - upsertUpload(updated); - } catch { - // Leave it pending. - } + // Phase 1: e621 MD5 lookup, 75 md5: metatags per query. + setPipeline({ phase: "md5", done: 0, total: targets.length }); + try { + for (let index = 0; index < targets.length; index += MD5_BATCH_SIZE) { + const md5Batch = targets.slice(index, index + MD5_BATCH_SIZE); + const posts = await fetchPostsByMd5( + creds, + md5Batch.map((temp) => temp.md5), + ); + for (const post of posts) { + const match = md5Batch.find((temp) => temp.md5 === post.file.md5); + if (!match) continue; + try { + const updated = await api( + `/api/uploads/${match.temp_id}/resolve/`, + { method: "POST", json: { mode: "link", post, auto: true } }, + ); + upsertUpload(updated); + } catch { + // Leave it pending. } } - } catch { - // e621 unavailable; leave everything pending. + setPipeline((current) => + current + ? { + ...current, + done: Math.min(index + MD5_BATCH_SIZE, targets.length), + } + : current, + ); } - invalidateUploads(); + } catch { + // e621 unavailable; those files stay pending. } + setPipeline(null); - // 2) IQDB similarity for every unresolved file that has not been checked - // yet (pending and local visual matches alike; IQDB is image-only). + // Phase 2: local visual similarity, one file at a time. + const cached = queryClient.getQueryData(["uploads"]) ?? []; + const visualTargets = cached.filter( + (temp) => targetIds.has(temp.temp_id) && temp.status === "pending", + ); + setPipeline({ phase: "visual", done: 0, total: visualTargets.length }); + for (const temp of visualTargets) { + try { + const updated = await api( + `/api/uploads/${temp.temp_id}/visual-match/`, + { method: "POST" }, + ); + upsertUpload(updated); + } catch { + // Keep it pending. + } + setPipeline((current) => + current ? { ...current, done: current.done + 1 } : current, + ); + } + setPipeline(null); + + // Phase 3: IQDB for everything still unresolved (serial queue with its + // own progress, cooldown and fatal-error stop). const fresh = await refreshUploads(); - void checkSimilarity(unresolvedIqdbIds(fresh)); + checkSimilarity(unresolvedIqdbIds(fresh)); } /** Mark an entry done, fade it out and drop it once it has been seen. */ @@ -1032,7 +1098,7 @@ export default function UploadPage() { invalidateUploads(); if (created.length > 0) { setProcessing(true); - void autoMatchAndCheck(created).finally(() => setProcessing(false)); + void runUploadPipeline(created).finally(() => setProcessing(false)); } } @@ -1133,33 +1199,38 @@ export default function UploadPage() { retry checks
- ) : busy || processing ? ( -
-

- - - {busy - ? `Uploading ${Math.min(settledCount + 1, batch.total)}/${batch.total}…` - : iqdbProgress - ? `Checking IQDB — ${iqdbProgress.done}/${iqdbProgress.total}${ - currentCheck ? ` · ${currentCheck}` : "" - }` - : "Processing uploads…"} - -

- {!busy && iqdbProgress ? ( -
-
-
- ) : null} -
+ ) : busy ? ( + + ) : pipeline ? ( + + ) : processing && iqdbProgress ? ( + + ) : processing ? ( +

+ Processing uploads… +

) : null}