JakeBreath
e2bf1c457f
Uploads were doing md5 + local visual matching inside the upload request (backend create) while the frontend later ran its own e621 MD5 pass, so the pipeline looked interleaved per file. Now every step is a phase applied to the whole batch in order: 1. upload (fast: md5 + exact-duplicate check only), 2. e621 MD5 lookup, 75 md5: metatags per posts.json request, 3. local visual similarity, one file at a time via the new POST /api/uploads/<id>/visual-match/ action, 4. IQDB through the existing serial queue. The board shows the active phase with its own progress bar (e621 MD5 in peach, visual in lavender, IQDB in teal) and every step updates the staged list as it lands. Verified from a headless run: one batched posts.json request for 10 files, then 10 visual-match calls, then IQDB.
413 lines
16 KiB
Python
413 lines
16 KiB
Python
"""Staged uploads: files land in a temp folder until they are resolved.
|
|
|
|
Flow:
|
|
- POST /api/uploads/ stage a file (computes MD5, detects duplicates)
|
|
- GET /api/uploads/ list my staged uploads
|
|
- GET /api/uploads/<id>/file/ serve the staged file (previews, IQDB)
|
|
- POST /api/uploads/<id>/iqdb/ store IQDB candidates -> visual_match
|
|
- POST /api/uploads/<id>/resolve/ link to a post or apply custom metadata
|
|
- DELETE /api/uploads/<id>/ discard a staged upload / dismiss a record
|
|
"""
|
|
|
|
import logging
|
|
import shutil
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
from django.conf import settings
|
|
from django.contrib.auth import get_user_model
|
|
from django.core import signing
|
|
from django.core.exceptions import ValidationError
|
|
from django.http import Http404
|
|
from django.utils import timezone
|
|
from rest_framework import mixins, status, viewsets
|
|
from rest_framework.decorators import action
|
|
from rest_framework.parsers import FormParser, JSONParser, MultiPartParser
|
|
from rest_framework.permissions import AllowAny
|
|
from rest_framework.response import Response
|
|
|
|
from . import services
|
|
from .models import MediaItem, TempUpload
|
|
from .permissions import CanUpload
|
|
from .serializers import TempUploadSerializer
|
|
from .tools import HASH_FIELDS, hashes_similarity
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def find_library_matches(path, limit=10, user=None, request=None):
|
|
"""Library items visually similar to a staged file."""
|
|
hashes = services.compute_visual_hashes(path)
|
|
if not hashes:
|
|
return []
|
|
algorithms = list(HASH_FIELDS)
|
|
threshold = settings.VISUAL_MATCH_THRESHOLD
|
|
matches = []
|
|
for item in MediaItem.objects.prefetch_related("locations"):
|
|
similarity = hashes_similarity(
|
|
hashes,
|
|
{field: getattr(item, field, "") for field in algorithms},
|
|
algorithms,
|
|
threshold,
|
|
)
|
|
if similarity is None:
|
|
continue
|
|
location = item.locations.first()
|
|
matches.append(
|
|
{
|
|
"j_id": f"J-{item.id}",
|
|
"filename": Path(location.rel_path).name if location else item.md5,
|
|
"similarity": round(similarity * 100, 1),
|
|
"thumbnail_url": services.signed_media_url(
|
|
item, user, "thumbnail", request=request
|
|
),
|
|
}
|
|
)
|
|
matches.sort(key=lambda entry: entry["similarity"], reverse=True)
|
|
return matches[:limit]
|
|
|
|
|
|
def complete_temp_upload(temp, download_url=None):
|
|
"""Index the upload into the library.
|
|
|
|
With a ``download_url`` the e621 original is fetched and the staged copy
|
|
is dropped; otherwise the staged file itself moves into the library.
|
|
"""
|
|
folder = Path(settings.WATCHED_FOLDER)
|
|
|
|
if download_url:
|
|
file_data = {}
|
|
if isinstance(temp.e621_data, dict):
|
|
file_data = temp.e621_data.get("file") or {}
|
|
extension = Path(urlparse(download_url).path).suffix
|
|
if not extension:
|
|
extension = Path(temp.original_filename).suffix
|
|
md5 = file_data.get("md5")
|
|
name = f"{md5}{extension}" if md5 else (temp.original_filename or "download")
|
|
destination = services.unique_destination(folder, name)
|
|
try:
|
|
services.download_file(download_url, destination)
|
|
except Exception:
|
|
destination.unlink(missing_ok=True)
|
|
raise
|
|
item, _, location, _ = services.index_file(destination, folder)
|
|
services.rename_location_to_j_id(item, location)
|
|
services.ensure_visual_hashes(item)
|
|
if temp.file:
|
|
temp.file.delete(save=False)
|
|
else:
|
|
if not temp.file:
|
|
raise ValueError("staged file is missing")
|
|
destination = services.unique_destination(
|
|
folder, temp.original_filename or "upload"
|
|
)
|
|
with temp.file.open("rb") as source, open(destination, "wb") as target:
|
|
shutil.copyfileobj(source, target)
|
|
item, _, location, _ = services.index_file(destination, folder)
|
|
services.rename_location_to_j_id(item, location)
|
|
services.ensure_visual_hashes(item)
|
|
temp.file.delete(save=False)
|
|
|
|
temp.library_item = item
|
|
temp.status = TempUpload.STATUS_COMPLETED
|
|
|
|
update_fields = []
|
|
if item.uploaded_by_id is None:
|
|
item.uploaded_by = temp.user
|
|
update_fields.append("uploaded_by")
|
|
if temp.e621_post_id is not None:
|
|
item.e621_post_id = temp.e621_post_id
|
|
item.e621_data = temp.e621_data
|
|
item.e621_match_status = MediaItem.E621_MATCHED
|
|
item.e621_checked_at = timezone.now()
|
|
update_fields += [
|
|
"e621_post_id",
|
|
"e621_data",
|
|
"e621_match_status",
|
|
"e621_checked_at",
|
|
]
|
|
if temp.custom_rating in {"s", "q", "e"} and not item.rating:
|
|
item.rating = temp.custom_rating
|
|
update_fields.append("rating")
|
|
if temp.custom_tags:
|
|
item.tags = temp.custom_tags
|
|
update_fields.append("tags")
|
|
if temp.custom_notes:
|
|
item.notes = temp.custom_notes
|
|
update_fields.append("notes")
|
|
if update_fields:
|
|
item.save(update_fields=update_fields + ["updated_at"])
|
|
|
|
temp.save()
|
|
return item
|
|
|
|
|
|
class TempUploadViewSet(
|
|
mixins.ListModelMixin,
|
|
mixins.RetrieveModelMixin,
|
|
mixins.DestroyModelMixin,
|
|
viewsets.GenericViewSet,
|
|
):
|
|
serializer_class = TempUploadSerializer
|
|
permission_classes = [CanUpload]
|
|
parser_classes = [MultiPartParser, FormParser, JSONParser]
|
|
# Unpaginated: the board shows every staged upload (69-file batches were
|
|
# silently cut to the API's 48-item page).
|
|
pagination_class = None
|
|
http_method_names = ["get", "post", "delete", "head", "options"]
|
|
|
|
def get_queryset(self):
|
|
queryset = TempUpload.objects.select_related("library_item")
|
|
user = self.request.user
|
|
if not user.is_app_staff:
|
|
queryset = queryset.filter(user=user)
|
|
return queryset
|
|
|
|
def create(self, request):
|
|
upload = request.FILES.get("file")
|
|
if upload is None:
|
|
return Response(
|
|
{"detail": "A file is required."}, status=status.HTTP_400_BAD_REQUEST
|
|
)
|
|
extension = Path(upload.name).suffix.lower()
|
|
if extension not in services.ALLOWED_EXTENSIONS:
|
|
return Response(
|
|
{"detail": f"Unsupported file type: {extension or 'unknown'}"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
temp = TempUpload.objects.create(
|
|
user=request.user,
|
|
file=upload,
|
|
original_filename=upload.name,
|
|
size=upload.size,
|
|
)
|
|
temp.md5 = services.compute_md5(temp.file.path)
|
|
|
|
existing = MediaItem.objects.filter(md5=temp.md5).first()
|
|
if existing is not None:
|
|
temp.status = TempUpload.STATUS_COMPLETED
|
|
temp.resolution = TempUpload.RESOLUTION_DUPLICATE
|
|
temp.library_item = existing
|
|
temp.file.delete(save=False)
|
|
# 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."""
|
|
user = request.user if request.user.is_authenticated else None
|
|
if user is None:
|
|
signature = request.query_params.get("sig")
|
|
if signature:
|
|
try:
|
|
payload = signing.loads(
|
|
signature,
|
|
salt=services.UPLOAD_FILE_SALT,
|
|
max_age=86400,
|
|
)
|
|
except signing.BadSignature:
|
|
payload = None
|
|
if payload and str(payload.get("temp")) == str(pk):
|
|
user = (
|
|
get_user_model()
|
|
.objects.filter(pk=payload.get("user"))
|
|
.first()
|
|
)
|
|
if user is None or not user.can_upload:
|
|
return Response(
|
|
{"detail": "Authentication required."},
|
|
status=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
|
|
temp = TempUpload.objects.filter(pk=pk).first()
|
|
is_owner = temp is not None and temp.user_id == user.id
|
|
if temp is None or not (is_owner or user.is_app_staff):
|
|
raise Http404
|
|
if not temp.file:
|
|
raise Http404
|
|
return services.serve_file(request, temp.file.path)
|
|
|
|
@action(detail=True, methods=["post"])
|
|
def iqdb(self, request, pk=None):
|
|
temp = self.get_object()
|
|
results = request.data.get("results")
|
|
if not isinstance(results, list):
|
|
return Response(
|
|
{"detail": "results must be a list."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
temp.iqdb_data = services.sanitize_iqdb_results(results)
|
|
if temp.status == TempUpload.STATUS_PENDING and temp.iqdb_data:
|
|
temp.status = TempUpload.STATUS_VISUAL_MATCH
|
|
temp.save(update_fields=["iqdb_data", "status", "updated_at"])
|
|
return Response(self.get_serializer(temp).data)
|
|
|
|
@action(detail=True, methods=["post"])
|
|
def resolve(self, request, pk=None):
|
|
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,
|
|
)
|
|
|
|
mode = str(request.data.get("mode") or "").strip()
|
|
download_url = None
|
|
if mode == "link":
|
|
post = request.data.get("post")
|
|
post_id = request.data.get("post_id")
|
|
trimmed = services.trim_e621_post(post)
|
|
has_numeric_id = post_id is not None and str(post_id).isdigit()
|
|
if trimmed is None and not has_numeric_id:
|
|
return Response(
|
|
{"detail": "A post payload or post_id is required."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
if trimmed is not None and trimmed.get("id") is not None:
|
|
temp.e621_post_id = int(trimmed["id"])
|
|
elif has_numeric_id:
|
|
temp.e621_post_id = int(post_id)
|
|
temp.e621_data = trimmed
|
|
temp.resolution = (
|
|
TempUpload.RESOLUTION_AUTO_MD5
|
|
if request.data.get("auto")
|
|
else TempUpload.RESOLUTION_LINKED
|
|
)
|
|
rating = str(request.data.get("rating") or "")
|
|
if rating in {"s", "q", "e"}:
|
|
temp.custom_rating = rating
|
|
tags = services.parse_tags(request.data.get("tags"))
|
|
if tags is not None:
|
|
temp.custom_tags = tags
|
|
notes = request.data.get("notes")
|
|
if notes is not None:
|
|
temp.custom_notes = str(notes)
|
|
|
|
# Prefer the e621 original; fall back to the staged file when the
|
|
# post has no URL or its file is byte-identical to the upload.
|
|
file_data = (trimmed or {}).get("file") or {}
|
|
candidate_url = str(
|
|
request.data.get("file_url") or file_data.get("url") or ""
|
|
).strip()
|
|
post_md5 = str(file_data.get("md5") or "").strip().lower()
|
|
if candidate_url and post_md5 and post_md5 == temp.md5.lower():
|
|
candidate_url = ""
|
|
download_url = candidate_url or None
|
|
elif mode == "custom":
|
|
rating = str(request.data.get("rating") or "")
|
|
temp.custom_rating = rating if rating in {"s", "q", "e"} else ""
|
|
temp.custom_tags = services.parse_tags(request.data.get("tags")) or []
|
|
temp.custom_notes = str(request.data.get("notes") or "")
|
|
temp.resolution = TempUpload.RESOLUTION_CUSTOM
|
|
else:
|
|
return Response(
|
|
{"detail": "mode must be 'link' or 'custom'."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
temp.save()
|
|
try:
|
|
complete_temp_upload(temp, download_url=download_url)
|
|
except Exception as exc: # noqa: BLE001 - report completion failures
|
|
logger.exception("Could not complete staged upload %s", temp.id)
|
|
temp.status = TempUpload.STATUS_ERROR
|
|
temp.save(update_fields=["status", "updated_at"])
|
|
return Response(
|
|
{"detail": f"Could not finish the upload: {exc}"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
temp.refresh_from_db()
|
|
return Response(self.get_serializer(temp).data)
|
|
|
|
@action(detail=False, methods=["post"], url_path="resolve-bulk")
|
|
def resolve_bulk(self, request):
|
|
"""Move many staged uploads into the library with one rating.
|
|
|
|
The upload board's bulk tool: pick a rating, tick pending files, and
|
|
they are all resolved as custom entries (their staged tags/notes are
|
|
kept). Own rows only, like the rest of the viewset.
|
|
"""
|
|
ids = request.data.get("temp_ids")
|
|
if not isinstance(ids, list) or not ids:
|
|
return Response(
|
|
{"detail": "temp_ids must be a non-empty list."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
rating = str(request.data.get("rating") or "").strip()
|
|
if rating not in {"s", "q", "e"}:
|
|
return Response(
|
|
{"detail": "rating must be one of s, q or e."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
try:
|
|
queryset = self.get_queryset().filter(
|
|
pk__in=[str(value) for value in ids]
|
|
)
|
|
except (ValidationError, ValueError):
|
|
return Response(
|
|
{"detail": "One or more ids are not valid upload ids."},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
resolved: list[str] = []
|
|
errors: list[dict[str, str]] = []
|
|
for temp in queryset:
|
|
if temp.status == TempUpload.STATUS_COMPLETED:
|
|
errors.append(
|
|
{"temp_id": str(temp.id), "error": "already in the library"}
|
|
)
|
|
continue
|
|
temp.custom_rating = rating
|
|
temp.resolution = TempUpload.RESOLUTION_CUSTOM
|
|
temp.save(update_fields=["custom_rating", "resolution", "updated_at"])
|
|
try:
|
|
complete_temp_upload(temp)
|
|
except Exception as exc: # noqa: BLE001 - report per-file failures
|
|
logger.exception("Could not complete staged upload %s", temp.id)
|
|
temp.status = TempUpload.STATUS_ERROR
|
|
temp.save(update_fields=["status", "updated_at"])
|
|
errors.append({"temp_id": str(temp.id), "error": str(exc)})
|
|
continue
|
|
resolved.append(str(temp.id))
|
|
|
|
found = {str(temp.id) for temp in queryset}
|
|
for value in sorted({str(value) for value in ids} - found):
|
|
errors.append({"temp_id": value, "error": "not found"})
|
|
|
|
return Response({"resolved": resolved, "errors": errors})
|
|
|
|
def perform_destroy(self, instance):
|
|
if instance.file:
|
|
instance.file.delete(save=False)
|
|
instance.delete()
|