Upload updates
This commit is contained in:
+182
-22
@@ -29,27 +29,32 @@ 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
|
||||
from .serializers import TempUploadListSerializer, TempUploadSerializer
|
||||
from .tools import HASH_FIELDS, hashed_items, 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 []
|
||||
def build_hash_index():
|
||||
"""Library hash mappings for similarity scans, loaded once per batch.
|
||||
|
||||
Only items that actually carry perceptual hashes are included; the old
|
||||
per-file scan walked every row (including videos and unchecked items).
|
||||
"""
|
||||
algorithms = list(HASH_FIELDS)
|
||||
return [
|
||||
(item, {field: getattr(item, field, "") for field in algorithms})
|
||||
for item in hashed_items(algorithms)
|
||||
]
|
||||
|
||||
|
||||
def match_hashes(hashes, index, limit=10, user=None, request=None):
|
||||
"""Library items whose perceptual hashes are close to ``hashes``."""
|
||||
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,
|
||||
)
|
||||
for item, item_hashes in index:
|
||||
similarity = hashes_similarity(hashes, item_hashes, algorithms, threshold)
|
||||
if similarity is None:
|
||||
continue
|
||||
location = item.locations.first()
|
||||
@@ -67,6 +72,20 @@ def find_library_matches(path, limit=10, user=None, request=None):
|
||||
return matches[:limit]
|
||||
|
||||
|
||||
def find_library_matches(path, limit=10, user=None, request=None, index=None):
|
||||
"""Library items visually similar to a staged file.
|
||||
|
||||
Pass a prebuilt ``index`` (see build_hash_index) to reuse it across a
|
||||
whole batch instead of rescanning the library per file.
|
||||
"""
|
||||
hashes = services.compute_visual_hashes(path)
|
||||
if not hashes:
|
||||
return []
|
||||
if index is None:
|
||||
index = build_hash_index()
|
||||
return match_hashes(hashes, index, limit=limit, user=user, request=request)
|
||||
|
||||
|
||||
def complete_temp_upload(temp, download_url=None):
|
||||
"""Index the upload into the library.
|
||||
|
||||
@@ -156,12 +175,20 @@ class TempUploadViewSet(
|
||||
pagination_class = None
|
||||
http_method_names = ["get", "post", "delete", "head", "options"]
|
||||
|
||||
def get_serializer_class(self):
|
||||
# The board polls the list, so its payload stays small; the metadata
|
||||
# modal fetches the full row from the detail endpoint.
|
||||
if self.action == "list":
|
||||
return TempUploadListSerializer
|
||||
return TempUploadSerializer
|
||||
|
||||
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
|
||||
# Staged uploads are private: everyone, staff included, only sees
|
||||
# their own board. (The file action still lets staff read bytes by id
|
||||
# for support purposes.)
|
||||
return TempUpload.objects.select_related("library_item").filter(
|
||||
user=self.request.user
|
||||
)
|
||||
|
||||
def create(self, request):
|
||||
upload = request.FILES.get("file")
|
||||
@@ -190,10 +217,14 @@ class TempUploadViewSet(
|
||||
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.
|
||||
# Visual similarity and IQDB run in the background pipeline so a large
|
||||
# batch uploads at full speed and the work survives the browser.
|
||||
temp.save()
|
||||
# Kick the server-side pipeline; staging no longer waits on e621 and
|
||||
# the work continues even if the browser navigates away.
|
||||
from .upload_pipeline import start_pipeline
|
||||
|
||||
start_pipeline(request.user)
|
||||
return Response(
|
||||
self.get_serializer(temp).data, status=status.HTTP_201_CREATED
|
||||
)
|
||||
@@ -221,6 +252,135 @@ class TempUploadViewSet(
|
||||
temp.save(update_fields=["visual_matches", "status", "updated_at"])
|
||||
return Response(self.get_serializer(temp).data)
|
||||
|
||||
@action(detail=False, methods=["get"])
|
||||
def status(self, request):
|
||||
"""Cheap pipeline state for the shell indicator and the upload page."""
|
||||
from .upload_pipeline import status_payload
|
||||
|
||||
return Response(status_payload(request.user))
|
||||
|
||||
@action(detail=False, methods=["post"])
|
||||
def process(self, request):
|
||||
"""Start (or resume) the pipeline for the caller's staged uploads.
|
||||
|
||||
Idempotent: the client calls this after staging files, on page load
|
||||
and when a paused run should be retried.
|
||||
"""
|
||||
from .upload_pipeline import start_pipeline, status_payload
|
||||
|
||||
start_pipeline(request.user)
|
||||
return Response(status_payload(request.user))
|
||||
|
||||
@action(detail=True, methods=["post"])
|
||||
def retry(self, request, pk=None):
|
||||
"""Queue one staged upload for another pipeline pass."""
|
||||
from .upload_pipeline import start_pipeline, status_payload
|
||||
|
||||
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,
|
||||
)
|
||||
update = {
|
||||
"claimed_at": None,
|
||||
"attempts": 0,
|
||||
"pipeline_error": "",
|
||||
"updated_at": timezone.now(),
|
||||
}
|
||||
# An explicit phase re-runs that one check even if it already ran.
|
||||
phase = str(request.data.get("phase") or "").strip()
|
||||
if phase == "md5":
|
||||
update["e621_checked_at"] = None
|
||||
elif phase == "visual":
|
||||
update["visual_matches"] = None
|
||||
update["visual_checked_at"] = None
|
||||
elif phase == "iqdb":
|
||||
update["iqdb_data"] = None
|
||||
if temp.status == TempUpload.STATUS_ERROR:
|
||||
# A failed import has to go through the MD5 phase again so the
|
||||
# completion is retried; other errors only re-run missing phases.
|
||||
update["status"] = TempUpload.STATUS_PENDING
|
||||
update["e621_checked_at"] = None
|
||||
elif temp.status not in (
|
||||
TempUpload.STATUS_PENDING,
|
||||
TempUpload.STATUS_VISUAL_MATCH,
|
||||
):
|
||||
update["status"] = TempUpload.STATUS_PENDING
|
||||
TempUpload.objects.filter(pk=temp.pk).update(**update)
|
||||
start_pipeline(request.user)
|
||||
return Response(status_payload(request.user))
|
||||
|
||||
@action(detail=False, methods=["post"], url_path="retry-all")
|
||||
def retry_all(self, request):
|
||||
"""Queue every retryable staged upload for another pipeline pass."""
|
||||
from .upload_pipeline import start_pipeline, status_payload
|
||||
|
||||
now = timezone.now()
|
||||
retryable = self.get_queryset().filter(
|
||||
status__in=[
|
||||
TempUpload.STATUS_PENDING,
|
||||
TempUpload.STATUS_VISUAL_MATCH,
|
||||
TempUpload.STATUS_ERROR,
|
||||
]
|
||||
)
|
||||
retryable.exclude(file="").update(
|
||||
claimed_at=None,
|
||||
attempts=0,
|
||||
pipeline_error="",
|
||||
updated_at=now,
|
||||
)
|
||||
retryable.exclude(file="").filter(status=TempUpload.STATUS_ERROR).update(
|
||||
status=TempUpload.STATUS_PENDING,
|
||||
e621_checked_at=None,
|
||||
)
|
||||
start_pipeline(request.user)
|
||||
return Response(status_payload(request.user))
|
||||
|
||||
@action(detail=False, methods=["post"], url_path="discard-bulk")
|
||||
def discard_bulk(self, request):
|
||||
"""Discard many staged uploads in one request (the board's "all")."""
|
||||
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,
|
||||
)
|
||||
if len(ids) > 1000:
|
||||
return Response(
|
||||
{"detail": "Too many ids in one request (max 1000)."},
|
||||
status=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
values = list(dict.fromkeys(str(value) for value in ids))
|
||||
try:
|
||||
queryset = self.get_queryset().filter(pk__in=values)
|
||||
except (ValidationError, ValueError):
|
||||
return Response(
|
||||
{"detail": "One or more temp_ids are not valid upload ids."},
|
||||
status=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
by_id = {str(temp.pk): temp for temp in queryset}
|
||||
|
||||
discarded: list[str] = []
|
||||
errors: list[dict[str, str]] = []
|
||||
for value in values:
|
||||
temp = by_id.get(value)
|
||||
if temp is None:
|
||||
errors.append({"temp_id": value, "error": "not found"})
|
||||
continue
|
||||
try:
|
||||
self.perform_destroy(temp)
|
||||
discarded.append(value)
|
||||
except Exception as exc: # noqa: BLE001 - report per-file failures
|
||||
logger.exception("Could not discard staged upload %s", value)
|
||||
errors.append({"temp_id": value, "error": str(exc)})
|
||||
return Response({"discarded": discarded, "errors": errors})
|
||||
|
||||
@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."""
|
||||
|
||||
Reference in New Issue
Block a user