JakeBreath
d0e2901c92
Backend: - TempUpload model: staged files (pending / visual_match / completed / error) with resolution, e621 payload, custom metadata and IQDB data - Files land in a temp folder and only move into the watched library folder once resolved; duplicates resolve immediately without a copy - Endpoints: stage (multipart), list, retrieve, temp file, IQDB save, resolve (link to post or custom metadata), discard/dismiss - cleanup_temp_uploads command for old staged files - Replaces the old direct-to-library upload endpoint Frontend: - Upload page is now a three-column board (Pending & Unmatched / Visual Similarity Detected / Auto-uploaded & Indexed) - After upload: MD5s are batch-checked against e621 and matches auto-complete with full post metadata; remaining files run through IQDB and move to the similarity column when candidates exist - Metadata modal with IQDB candidates, post-ID linking and custom tags/rating/notes; discard and dismiss actions - e621 client gains fetchPostsByMd5 and iqdbSearch helpers Roadmap updated with the completed upload items.
278 lines
8.6 KiB
Python
278 lines
8.6 KiB
Python
import hashlib
|
|
import json
|
|
import mimetypes
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
from django.conf import settings
|
|
from django.http import FileResponse, Http404, HttpResponse
|
|
from django.utils.text import get_valid_filename
|
|
|
|
from .models import MediaItem, MediaLocation
|
|
|
|
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".apng", ".mp4", ".webm"}
|
|
VIDEO_EXTENSIONS = {".mp4", ".webm"}
|
|
CHUNK_SIZE = 1024 * 1024
|
|
RANGE_RE = re.compile(r"bytes=(\d*)-(\d*)$")
|
|
|
|
|
|
def compute_md5(path):
|
|
digest = hashlib.md5()
|
|
with open(path, "rb") as handle:
|
|
for chunk in iter(lambda: handle.read(CHUNK_SIZE), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def index_file(path, folder):
|
|
"""Index one file into MediaItem/MediaLocation.
|
|
|
|
Returns (item, created_item, created_location).
|
|
"""
|
|
path = Path(path)
|
|
folder = Path(folder)
|
|
stat = path.stat()
|
|
md5 = compute_md5(path)
|
|
item, created_item = MediaItem.objects.get_or_create(
|
|
md5=md5, defaults={"size": stat.st_size}
|
|
)
|
|
if not created_item and item.size != stat.st_size:
|
|
item.size = stat.st_size
|
|
item.save(update_fields=["size", "updated_at"])
|
|
location, created_location = MediaLocation.objects.update_or_create(
|
|
item=item,
|
|
path=str(path),
|
|
defaults={
|
|
"rel_path": str(path.relative_to(folder)),
|
|
"mtime": stat.st_mtime,
|
|
},
|
|
)
|
|
return item, created_item, created_location
|
|
|
|
|
|
def parse_tags(raw):
|
|
"""Normalize a comma-separated string or JSON list into a list of tags."""
|
|
if raw is None:
|
|
return None
|
|
if isinstance(raw, (list, tuple)):
|
|
values = raw
|
|
else:
|
|
text = str(raw).strip()
|
|
if not text:
|
|
return None
|
|
if text.startswith("["):
|
|
try:
|
|
values = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
values = text.split(",")
|
|
else:
|
|
values = text.split(",")
|
|
cleaned = []
|
|
for value in values:
|
|
tag = str(value).strip()
|
|
if tag and tag not in cleaned:
|
|
cleaned.append(tag[:100])
|
|
return cleaned
|
|
|
|
|
|
def unique_destination(folder, filename):
|
|
"""Return a non-existing path inside folder for the given filename."""
|
|
folder = Path(folder)
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
filename = get_valid_filename(Path(filename).name) or "upload"
|
|
candidate = folder / filename
|
|
stem, suffix = candidate.stem, candidate.suffix
|
|
counter = 1
|
|
while candidate.exists():
|
|
candidate = folder / f"{stem}-{counter}{suffix}"
|
|
counter += 1
|
|
return candidate
|
|
|
|
|
|
class RangeFileWrapper:
|
|
"""Iterate over a limited byte range of an open file."""
|
|
|
|
def __init__(self, file, length, chunk_size=64 * 1024):
|
|
self.file = file
|
|
self.remaining = length
|
|
self.chunk_size = chunk_size
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self.remaining <= 0:
|
|
raise StopIteration
|
|
data = self.file.read(min(self.chunk_size, self.remaining))
|
|
if not data:
|
|
raise StopIteration
|
|
self.remaining -= len(data)
|
|
return data
|
|
|
|
def close(self):
|
|
self.file.close()
|
|
|
|
|
|
def serve_file(request, path, download=False):
|
|
"""Serve a file with HTTP range support (needed for video seeking)."""
|
|
path = Path(path)
|
|
if not path.is_file():
|
|
raise Http404
|
|
size = path.stat().st_size
|
|
content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
|
|
range_header = request.headers.get("Range", "").strip()
|
|
if range_header:
|
|
match = RANGE_RE.match(range_header)
|
|
if match:
|
|
start_raw, end_raw = match.groups()
|
|
if start_raw == "" and end_raw:
|
|
length = min(int(end_raw), size)
|
|
start, end = size - length, size - 1
|
|
else:
|
|
start = int(start_raw or 0)
|
|
end = min(int(end_raw) if end_raw else size - 1, size - 1)
|
|
if start >= size or start > end:
|
|
response = HttpResponse(status=416)
|
|
response["Content-Range"] = f"bytes */{size}"
|
|
return response
|
|
length = end - start + 1
|
|
handle = open(path, "rb")
|
|
handle.seek(start)
|
|
response = FileResponse(
|
|
RangeFileWrapper(handle, length),
|
|
status=206,
|
|
content_type=content_type,
|
|
)
|
|
response["Content-Length"] = str(length)
|
|
response["Content-Range"] = f"bytes {start}-{end}/{size}"
|
|
response["Accept-Ranges"] = "bytes"
|
|
return response
|
|
response = FileResponse(
|
|
open(path, "rb"),
|
|
content_type=content_type,
|
|
as_attachment=download,
|
|
filename=path.name,
|
|
)
|
|
response["Accept-Ranges"] = "bytes"
|
|
return response
|
|
|
|
|
|
def download_file(url, destination):
|
|
"""Stream a remote file into destination (used by Download to Library)."""
|
|
import requests
|
|
|
|
headers = {"User-Agent": settings.USER_AGENT}
|
|
with requests.get(url, headers=headers, stream=True, timeout=120) as response:
|
|
response.raise_for_status()
|
|
with open(destination, "wb") as handle:
|
|
for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
|
|
if chunk:
|
|
handle.write(chunk)
|
|
|
|
|
|
E621_DESCRIPTION_LIMIT = 20000
|
|
|
|
|
|
def trim_e621_post(post):
|
|
"""Keep a compact, size-bounded copy of an e621 post payload."""
|
|
if not isinstance(post, dict):
|
|
return None
|
|
file_data = post.get("file") or {}
|
|
relationships = post.get("relationships") or {}
|
|
score = post.get("score") or {}
|
|
categories = post.get("tags") or {}
|
|
return {
|
|
"id": post.get("id"),
|
|
"created_at": post.get("created_at"),
|
|
"rating": post.get("rating"),
|
|
"tags": {
|
|
str(category): [str(tag)[:200] for tag in tags][:500]
|
|
for category, tags in categories.items()
|
|
if isinstance(tags, list)
|
|
},
|
|
"score": {
|
|
"up": score.get("up"),
|
|
"down": score.get("down"),
|
|
"total": score.get("total"),
|
|
},
|
|
"fav_count": post.get("fav_count"),
|
|
"comment_count": post.get("comment_count"),
|
|
"sources": [
|
|
str(source)[:500]
|
|
for source in (post.get("sources") or [])
|
|
if source
|
|
][:20],
|
|
"description": str(post.get("description") or "")[:E621_DESCRIPTION_LIMIT],
|
|
"pools": [
|
|
int(pool)
|
|
for pool in (post.get("pools") or [])
|
|
if str(pool).isdigit()
|
|
][:50],
|
|
"relationships": {
|
|
"parent_id": relationships.get("parent_id"),
|
|
"has_children": relationships.get("has_children"),
|
|
},
|
|
"file": {
|
|
"md5": file_data.get("md5"),
|
|
"ext": file_data.get("ext"),
|
|
"size": file_data.get("size"),
|
|
"width": file_data.get("width"),
|
|
"height": file_data.get("height"),
|
|
},
|
|
"uploader_name": post.get("uploader_name"),
|
|
}
|
|
|
|
|
|
def sanitize_iqdb_results(results):
|
|
"""Keep a compact, size-bounded copy of IQDB candidates."""
|
|
cleaned = []
|
|
for result in results[:10]:
|
|
if not isinstance(result, dict):
|
|
continue
|
|
score = result.get("score")
|
|
post_id = result.get("post_id")
|
|
entry = {
|
|
"post_id": post_id if isinstance(post_id, int) else None,
|
|
"score": float(score) if isinstance(score, (int, float)) else None,
|
|
"preview_url": str(result.get("preview_url") or "")[:500] or None,
|
|
"rating": str(result.get("rating") or "")[:1] or None,
|
|
"md5": str(result.get("md5") or "")[:32] or None,
|
|
}
|
|
if entry["post_id"] or entry["preview_url"]:
|
|
cleaned.append(entry)
|
|
return cleaned
|
|
|
|
|
|
def generate_video_thumbnail(md5, path):
|
|
"""Extract a JPEG thumbnail from a video, cached under MEDIA_ROOT/thumbs."""
|
|
if not shutil.which("ffmpeg"):
|
|
return None
|
|
thumbs_dir = Path(settings.MEDIA_ROOT) / "thumbs"
|
|
thumbs_dir.mkdir(parents=True, exist_ok=True)
|
|
target = thumbs_dir / f"{md5}.jpg"
|
|
if target.exists() and target.stat().st_mtime >= os.path.getmtime(path):
|
|
return target
|
|
command = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-ss",
|
|
"0.5",
|
|
"-i",
|
|
str(path),
|
|
"-frames:v",
|
|
"1",
|
|
"-vf",
|
|
"scale=480:-2",
|
|
"-loglevel",
|
|
"error",
|
|
str(target),
|
|
]
|
|
try:
|
|
subprocess.run(command, check=True, capture_output=True, timeout=60)
|
|
except (subprocess.SubprocessError, OSError):
|
|
return None
|
|
return target if target.exists() else None
|