Bryan Van Deusen
548 lines
19 KiB
Python
548 lines
19 KiB
Python
# app/utils/image_importer.py
|
||
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
import hashlib
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import tempfile
|
||
import time
|
||
import uuid
|
||
|
||
from PIL import Image
|
||
import exifread
|
||
import imagehash
|
||
|
||
from app import db
|
||
from app.models import ImageRecord, Tag, ArchiveRecord
|
||
|
||
|
||
# =============================================================================
|
||
# Configuration & Constants
|
||
# =============================================================================
|
||
|
||
THUMB_SIZE = (400, 400)
|
||
|
||
# Media and archive extensions we handle
|
||
ALLOWED_MEDIA_EXTS = (
|
||
".jpg", ".jpeg", ".jfif", ".png", ".gif", ".bmp", ".tiff", ".webp", ".mp4", ".mov"
|
||
)
|
||
ALLOWED_ARCHIVE_EXTS = (
|
||
".zip", ".rar", ".cbr", ".7z",
|
||
".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz"
|
||
)
|
||
|
||
# Regex for multi-part detection
|
||
RAR_PART_RE = re.compile(r"\.part(\d+)\.rar$", re.IGNORECASE)
|
||
SEVENZ_PART_RE = re.compile(r"\.7z\.(\d{3})$", re.IGNORECASE)
|
||
OLD_RAR_VOL_RE = re.compile(r"\.r\d{2}$", re.IGNORECASE) # .r00, .r01, ...
|
||
|
||
# Environment configuration (optional)
|
||
ENV_TMPDIR = "ARCHIVE_TMPDIR" # where to extract archives (default: system tmp)
|
||
ENV_MINFREE = "ARCHIVE_MIN_FREE_GB" # min free GB required to start extraction (float; 0 = disabled)
|
||
|
||
|
||
# =============================================================================
|
||
# Utilities: Temp space management & free-space guard
|
||
# =============================================================================
|
||
|
||
def get_tmp_base() -> Path:
|
||
"""Return the base temp directory for archive extraction."""
|
||
return Path(os.environ.get(ENV_TMPDIR, tempfile.gettempdir())).resolve()
|
||
|
||
|
||
def cleanup_stale_tempdirs(prefix: str = "extract_") -> None:
|
||
"""Delete any leftover extraction directories from previous crashes."""
|
||
base = get_tmp_base()
|
||
base.mkdir(parents=True, exist_ok=True)
|
||
for entry in base.glob(f"{prefix}*"):
|
||
if entry.is_dir():
|
||
try:
|
||
shutil.rmtree(entry, ignore_errors=True)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def ensure_free_space_or_raise(path: Path, min_free_gb: float) -> None:
|
||
"""Raise if there is not enough free space at path."""
|
||
if min_free_gb <= 0:
|
||
return
|
||
total, used, free = shutil.disk_usage(str(path))
|
||
free_gb = free / (1024 ** 3)
|
||
if free_gb < min_free_gb:
|
||
raise RuntimeError(
|
||
f"Insufficient free space at {path}: {free_gb:.1f} GB free "
|
||
f"(required >= {min_free_gb:.1f} GB)."
|
||
)
|
||
|
||
|
||
# =============================================================================
|
||
# Utilities: External extractors with resilient fallback
|
||
# =============================================================================
|
||
|
||
class ExtractError(Exception):
|
||
"""Raised when archive extraction fails with all strategies."""
|
||
|
||
|
||
def _run_cmd(cmd: list[str]) -> tuple[int, str, str]:
|
||
p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||
return p.returncode, p.stdout, p.stderr
|
||
|
||
|
||
def extract_with_unar(archive_path: str, out_dir: str) -> None:
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
rc, out, err = _run_cmd(["unar", "-o", out_dir, archive_path])
|
||
if rc != 0:
|
||
raise ExtractError(f"unar failed (rc={rc})\nSTDOUT:\n{out}\nSTDERR:\n{err}")
|
||
|
||
|
||
def extract_with_7z(archive_path: str, out_dir: str) -> None:
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
rc, out, err = _run_cmd(["7z", "x", "-y", f"-o{out_dir}", archive_path])
|
||
if rc != 0:
|
||
raise ExtractError(f"7z failed (rc={rc})\nSTDOUT:\n{out}\nSTDERR:\n{err}")
|
||
|
||
|
||
def looks_password_protected_with_lsar(archive_path: str) -> bool:
|
||
"""Best-effort pre-check for 'Encrypted/Password' in lsar output."""
|
||
try:
|
||
rc, out, err = _run_cmd(["lsar", archive_path])
|
||
txt = (out + "\n" + err).lower()
|
||
return ("encrypted" in txt) or ("password" in txt)
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def extract_archive_resilient(archive_path: str, out_dir: str) -> None:
|
||
"""
|
||
Prefer unar for RAR/CBR (RAR5-friendly); fallback to 7z.
|
||
For other formats, try 7z first then unar.
|
||
Surfaces stderr so logs show CRC/password/permissions issues.
|
||
"""
|
||
lower = archive_path.lower()
|
||
errors = []
|
||
|
||
# Optional password hint to skip early
|
||
if lower.endswith((".rar", ".cbr")) and looks_password_protected_with_lsar(archive_path):
|
||
raise ExtractError("Archive appears password-protected; skipping (no password provided).")
|
||
|
||
if lower.endswith((".rar", ".cbr")):
|
||
order = (extract_with_unar, extract_with_7z)
|
||
else:
|
||
order = (extract_with_7z, extract_with_unar)
|
||
|
||
for fn in order:
|
||
try:
|
||
fn(archive_path, out_dir)
|
||
return
|
||
except ExtractError as e:
|
||
errors.append(str(e))
|
||
|
||
raise ExtractError("Both extractors failed:\n" + "\n---\n".join(errors))
|
||
|
||
|
||
# =============================================================================
|
||
# Core: Import task & archive processing
|
||
# =============================================================================
|
||
|
||
def import_images_task(source_dir: str, dest_dir: str) -> str:
|
||
"""
|
||
Walk the import tree, process archives first (only first volumes),
|
||
then media files. Batches DB commits for efficiency.
|
||
"""
|
||
# Cleanup any stale extraction dirs from previous runs/crashes
|
||
cleanup_stale_tempdirs()
|
||
|
||
imported: list[str] = []
|
||
batch_size = 10
|
||
batch_counter = 0
|
||
|
||
# Preload existing pHashes to avoid N+1 lookups
|
||
existing_phashes = [
|
||
(imagehash.hex_to_hash(img.perceptual_hash), img.width, img.height)
|
||
for img in ImageRecord.query.filter(ImageRecord.perceptual_hash != None).all()
|
||
]
|
||
|
||
# Ensure thumbs base exists
|
||
(Path(dest_dir) / "thumbs").mkdir(parents=True, exist_ok=True)
|
||
|
||
# Optional free-space guard
|
||
min_free = float(os.environ.get(ENV_MINFREE, "0") or "0")
|
||
try:
|
||
ensure_free_space_or_raise(get_tmp_base(), min_free)
|
||
except Exception as e:
|
||
print(f"[WARN] Free-space check failed: {e}")
|
||
|
||
for artist_dir in os.listdir(source_dir):
|
||
artist_path = os.path.join(source_dir, artist_dir)
|
||
if not os.path.isdir(artist_path):
|
||
continue
|
||
|
||
for root, _, files in os.walk(artist_path):
|
||
for file in files:
|
||
full_path = os.path.join(root, file)
|
||
lower = file.lower()
|
||
|
||
# Archives first
|
||
if lower.endswith(ALLOWED_ARCHIVE_EXTS):
|
||
if not is_first_volume(full_path):
|
||
print(f"[SKIP] Not first archive volume: {file}")
|
||
continue
|
||
try:
|
||
count = process_archive(
|
||
archive_path=full_path,
|
||
dest_dir=dest_dir,
|
||
artist=artist_dir,
|
||
existing_phashes=existing_phashes
|
||
)
|
||
imported.append(f"[archive]{file}:{count}")
|
||
except Exception as e:
|
||
print(f"[WARN] Failed processing archive {file}: {e}")
|
||
continue
|
||
|
||
# Non-archive media
|
||
if not lower.endswith(ALLOWED_MEDIA_EXTS):
|
||
continue
|
||
|
||
added, phash, _rec = import_single_file(
|
||
src_path=full_path,
|
||
dest_dir=dest_dir,
|
||
artist=artist_dir,
|
||
existing_phashes=existing_phashes
|
||
)
|
||
if added:
|
||
imported.append(file)
|
||
|
||
if phash:
|
||
md = extract_metadata(full_path)
|
||
if md.get("width") and md.get("height"):
|
||
existing_phashes.append((phash, md["width"], md["height"]))
|
||
|
||
batch_counter += 1
|
||
if batch_counter >= batch_size:
|
||
db.session.commit()
|
||
print(f"[INFO] Committed batch of {batch_size} items.")
|
||
batch_counter = 0
|
||
if "sqlite" in db.engine.url.drivername:
|
||
time.sleep(1)
|
||
|
||
if batch_counter > 0:
|
||
db.session.commit()
|
||
print(f"[INFO] Committed final batch of {batch_counter} items.")
|
||
|
||
summary = f"Imported {len(imported)} items."
|
||
print(f"[INFO] {summary}")
|
||
return summary
|
||
|
||
|
||
def process_archive(archive_path: str, dest_dir: str, artist: str, existing_phashes) -> int:
|
||
"""
|
||
Extract archive into a temp dir, import/tag media, and record ArchiveRecord for de-dup.
|
||
Returns number of items imported or tagged.
|
||
"""
|
||
archive_size = os.path.getsize(archive_path)
|
||
archive_hash = calculate_hash(archive_path)
|
||
|
||
# De-dup the archive itself
|
||
existing_archive = ArchiveRecord.query.filter_by(hash=archive_hash).first()
|
||
if existing_archive:
|
||
print(f"[SKIP] Archive already imported: {archive_path}")
|
||
return 0
|
||
|
||
# Create/reuse tag for this archive
|
||
tag_name = f"archive:{artist}/{os.path.basename(archive_path)}"
|
||
tag = Tag.query.filter_by(name=tag_name).first()
|
||
if not tag:
|
||
tag = Tag(name=tag_name, kind="archive")
|
||
db.session.add(tag)
|
||
db.session.flush() # ensure tag.id
|
||
|
||
# Create an extraction dir under the configured temp base
|
||
tmpdir = tempfile.mkdtemp(prefix="extract_", dir=str(get_tmp_base()))
|
||
imported_or_tagged = 0
|
||
|
||
try:
|
||
# Resilient extraction (unar → fallback 7z for RAR; 7z → fallback unar for others)
|
||
extract_archive_resilient(archive_path, tmpdir)
|
||
|
||
# Walk extracted tree and process media files
|
||
for root, _, files in os.walk(tmpdir):
|
||
for file in files:
|
||
if not file.lower().endswith(ALLOWED_MEDIA_EXTS):
|
||
continue
|
||
src_path = os.path.join(root, file)
|
||
|
||
added, _phash, rec = import_single_file(
|
||
src_path=src_path,
|
||
dest_dir=dest_dir,
|
||
artist=artist,
|
||
existing_phashes=existing_phashes,
|
||
commit=False # batch after loop
|
||
)
|
||
if rec is not None and tag not in rec.tags:
|
||
rec.tags.append(tag)
|
||
imported_or_tagged += 1
|
||
|
||
# Record this archive as processed
|
||
arch = ArchiveRecord(
|
||
filename=archive_path,
|
||
file_size=archive_size,
|
||
hash=archive_hash,
|
||
artist=artist,
|
||
tag=tag
|
||
)
|
||
db.session.add(arch)
|
||
db.session.commit()
|
||
print(f"[INFO] Archive processed: {archive_path} items={imported_or_tagged}")
|
||
return imported_or_tagged
|
||
|
||
finally:
|
||
shutil.rmtree(tmpdir, ignore_errors=True)
|
||
|
||
|
||
# =============================================================================
|
||
# Core: Single-file import path
|
||
# =============================================================================
|
||
|
||
def import_single_file(src_path: str, dest_dir: str, artist: str, existing_phashes, commit: bool = True):
|
||
"""
|
||
Import a single media file. If an equivalent file already exists, return that
|
||
record so callers can tag it (no duplicate import).
|
||
Returns (added: bool, phash_or_None, record_or_None).
|
||
"""
|
||
filename = os.path.basename(src_path)
|
||
print(f"[INFO] Processing: {src_path}")
|
||
file_size = os.path.getsize(src_path)
|
||
|
||
# Check by name + size
|
||
existing = ImageRecord.query.filter_by(filename=filename, file_size=file_size).first()
|
||
if existing:
|
||
print(f"[SKIP] Duplicate by name+size: {filename}")
|
||
return (False, None, existing)
|
||
|
||
# Check by content hash
|
||
file_hash = calculate_hash(src_path)
|
||
existing = ImageRecord.query.filter_by(hash=file_hash).first()
|
||
if existing:
|
||
print(f"[SKIP] Duplicate by hash: {filename}")
|
||
return (False, None, existing)
|
||
|
||
# Metadata & pHash similarity
|
||
metadata = extract_metadata(src_path)
|
||
if not metadata["width"] or not metadata["height"]:
|
||
print(f"[SKIP] Missing dimension data for {filename}")
|
||
return (False, None, None)
|
||
|
||
phash = calculate_perceptual_hash(src_path)
|
||
if phash and is_similar_image(phash, metadata["width"], metadata["height"], existing_phashes):
|
||
print(f"[SKIP] {filename} is visually similar to an existing larger image.")
|
||
return (False, phash, None)
|
||
|
||
# Copy into /images/<artist>/<filename>
|
||
target_dir = os.path.join(dest_dir, artist)
|
||
os.makedirs(target_dir, exist_ok=True)
|
||
dest_path = os.path.join(target_dir, filename)
|
||
shutil.copy2(src_path, dest_path)
|
||
|
||
# Thumbnail
|
||
try:
|
||
if dest_path.lower().endswith((".mp4", ".mov")):
|
||
thumb_filename = f"{uuid.uuid4().hex}.jpg"
|
||
thumb_path = os.path.join(dest_dir, "thumbs", thumb_filename)
|
||
thumb_path = generate_video_thumbnail(dest_path, thumb_path)
|
||
else:
|
||
thumb_path = generate_thumbnail(dest_path)
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to generate thumbnail for {filename}: {e}")
|
||
thumb_path = None
|
||
|
||
# Create DB record
|
||
record = ImageRecord(
|
||
filename=filename,
|
||
filepath=dest_path,
|
||
thumb_path=thumb_path,
|
||
hash=file_hash,
|
||
perceptual_hash=str(phash) if phash else None,
|
||
artist=artist,
|
||
file_size=metadata["file_size"],
|
||
width=metadata["width"],
|
||
height=metadata["height"],
|
||
format=metadata["format"],
|
||
camera_model=metadata["camera_model"],
|
||
taken_at=metadata["taken_at"],
|
||
imported_at=datetime.utcnow()
|
||
)
|
||
db.session.add(record)
|
||
if commit:
|
||
db.session.commit()
|
||
return (True, phash, record)
|
||
|
||
|
||
# =============================================================================
|
||
# Helpers: Thumbnails, metadata, hashing, similarity
|
||
# =============================================================================
|
||
|
||
def generate_thumbnail(image_path: str, size: tuple[int, int] = (400, 400), overwrite: bool = False) -> str:
|
||
"""
|
||
Save thumbnails mirrored under /images/thumbs/<artist>/<filename>.
|
||
"""
|
||
images_root = Path("/images").resolve()
|
||
image_path = Path(image_path).resolve()
|
||
rel_path = image_path.relative_to(images_root) # raises if not under /images
|
||
|
||
thumb_path = images_root / "thumbs" / rel_path
|
||
thumb_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
if not overwrite and thumb_path.exists():
|
||
return str(thumb_path)
|
||
|
||
with Image.open(image_path) as img:
|
||
img.thumbnail(size)
|
||
img.save(thumb_path)
|
||
|
||
return str(thumb_path)
|
||
|
||
|
||
def generate_video_thumbnail(video_path: str, output_path: str, time_position: str = "00:00:01") -> str | None:
|
||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||
command = [
|
||
"ffmpeg", "-i", video_path, "-ss", time_position,
|
||
"-vframes", "1", "-vf", f"scale={THUMB_SIZE[0]}:-1", output_path
|
||
]
|
||
try:
|
||
subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||
return output_path
|
||
except subprocess.CalledProcessError as e:
|
||
print(f"[WARN] Failed to extract video thumbnail for {video_path}: {e}")
|
||
return None
|
||
|
||
|
||
def calculate_hash(file_path: str) -> str:
|
||
hash_func = hashlib.sha256()
|
||
with open(file_path, "rb") as f:
|
||
for chunk in iter(lambda: f.read(1024 * 1024), b""): # 1MB chunks work well for big files
|
||
hash_func.update(chunk)
|
||
return hash_func.hexdigest()
|
||
|
||
|
||
def calculate_perceptual_hash(file_path: str):
|
||
try:
|
||
with Image.open(file_path) as img:
|
||
return imagehash.phash(img)
|
||
except Exception:
|
||
# videos or unreadable images land here (phash not applicable)
|
||
return None
|
||
|
||
|
||
def is_similar_image(phash, width: int, height: int, existing_phashes, threshold: int = 2) -> bool:
|
||
for existing, ew, eh in existing_phashes:
|
||
try:
|
||
distance = phash - existing
|
||
if distance <= threshold and ew >= width and eh >= height:
|
||
return True
|
||
except Exception as e:
|
||
print(f"[WARN] Failed pHash comparison: {e}")
|
||
return False
|
||
|
||
|
||
def extract_metadata(file_path: str) -> dict:
|
||
metadata = {
|
||
"file_size": None,
|
||
"width": None,
|
||
"height": None,
|
||
"format": None,
|
||
"camera_model": None,
|
||
"taken_at": None
|
||
}
|
||
|
||
metadata["file_size"] = os.path.getsize(file_path)
|
||
ext = os.path.splitext(file_path)[1].lower()
|
||
|
||
# Video: use ffprobe
|
||
if ext in (".mp4", ".mov"):
|
||
try:
|
||
cmd = [
|
||
"ffprobe", "-v", "error",
|
||
"-select_streams", "v:0",
|
||
"-show_entries", "stream=width,height",
|
||
"-of", "json", file_path
|
||
]
|
||
result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, check=False)
|
||
data = json.loads(result.stdout or "{}")
|
||
stream = data.get("streams", [{}])[0]
|
||
metadata["width"] = stream.get("width")
|
||
metadata["height"] = stream.get("height")
|
||
metadata["format"] = ext.lstrip(".")
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to extract video metadata: {file_path} – {e}")
|
||
return metadata
|
||
|
||
# Image: PIL + EXIF (best-effort)
|
||
try:
|
||
with Image.open(file_path) as img:
|
||
metadata["width"], metadata["height"] = img.size
|
||
metadata["format"] = (img.format or "").lower()
|
||
except Exception as e:
|
||
mime_type, _ = mimetypes.guess_type(file_path)
|
||
print(f"[WARN] Failed to read image metadata: {file_path} ({e}) ext={ext} mime={mime_type}")
|
||
|
||
try:
|
||
with open(file_path, "rb") as f:
|
||
tags = exifread.process_file(f, stop_tag="UNDEF", details=False)
|
||
if "EXIF DateTimeOriginal" in tags:
|
||
metadata["taken_at"] = datetime.strptime(str(tags["EXIF DateTimeOriginal"]), "%Y:%m:%d %H:%M:%S")
|
||
if "Image Model" in tags:
|
||
metadata["camera_model"] = str(tags["Image Model"])
|
||
except Exception:
|
||
pass
|
||
|
||
return metadata
|
||
|
||
|
||
# =============================================================================
|
||
# Helpers: Multi-part detection
|
||
# =============================================================================
|
||
|
||
def is_first_volume(path: str) -> bool:
|
||
"""
|
||
Determine if 'path' is the first volume of a multi-part archive.
|
||
We only attempt extraction from the first volume.
|
||
"""
|
||
name = os.path.basename(path).lower()
|
||
|
||
# RAR: .partN.rar (only N==1 is first)
|
||
m = RAR_PART_RE.search(name)
|
||
if m:
|
||
n = m.group(1)
|
||
return n in ("1", "01", "001")
|
||
|
||
# RAR old-style: first is .rar, subsequent are .r00, .r01...
|
||
if name.endswith(".rar"):
|
||
return True
|
||
if OLD_RAR_VOL_RE.search(name):
|
||
return False
|
||
|
||
# CBR is typically a single-volume RAR
|
||
if name.endswith(".cbr"):
|
||
return True
|
||
|
||
# 7z multi-part: .7z.001 is first
|
||
m = SEVENZ_PART_RE.search(name)
|
||
if m:
|
||
n = m.group(1)
|
||
return n in ("1", "01", "001")
|
||
if name.endswith(".7z"):
|
||
return True # single-volume 7z
|
||
|
||
# ZIP and tarballs: treat as single-volume (we don't include .z01 in our allowlist)
|
||
if name.endswith((".zip", ".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz")):
|
||
return True
|
||
|
||
return True
|