This repository has been archived on 2026-05-31. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
imagerepo/app/utils/image_importer.py
T
2025-08-16 10:25:16 -04:00

752 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 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, ImageOps
import exifread
import imagehash
from pyunpack import Archive
from app import db
from app.models import ImageRecord, Tag, ArchiveRecord
FFMPEG_THUMB_TIMEOUT = int(os.environ.get("THUMBS_FFMPEG_TIMEOUT", "60")) # seconds
Image.MAX_IMAGE_PIXELS = int(os.environ.get("PIL_MAX_IMAGE_PIXELS", "178956970"))
ARCHIVE_NUM_WIDTH = int(os.environ.get("ARCHIVE_NUM_WIDTH", "4"))
# =============================================================================
# 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, dest_dir, artist, existing_phashes):
"""
Extracts archive to temp dir, imports/tag existing images, and records ArchiveRecord.
Creates an archive:* tag ONLY if at least one image was imported or tagged.
Returns number of items imported or tagged.
"""
archive_size = os.path.getsize(archive_path)
archive_hash = calculate_hash(archive_path)
# Skip if this exact archive file was already processed (hash-based)
existing_archive = ArchiveRecord.query.filter_by(hash=archive_hash).first()
if existing_archive:
print(f"[SKIP] Archive already imported: {archive_path}")
return 0
tmpdir = tempfile.mkdtemp(prefix="extract_")
touched_records = [] # ImageRecord objects we imported or matched
imported_or_tagged = 0
created_tag = None
try:
# Extract using pyunpack
Archive(archive_path).extractall(tmpdir)
# Import / collect all media
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, # still used for auto artist:<name> tag
existing_phashes=existing_phashes,
commit=False
)
if rec is not None:
touched_records.append(rec)
if added:
imported_or_tagged += 1
# Only now: create/attach the archive tag if we actually touched images
if touched_records:
# Incremental name like 'archive:{artist}/0001'
tag_name = compute_next_archive_tag_name(artist)
created_tag = Tag.query.filter_by(name=tag_name).first()
if not created_tag:
created_tag = Tag(name=tag_name, kind='archive')
db.session.add(created_tag)
db.session.flush()
newly_tagged = 0
for rec in touched_records:
if created_tag not in rec.tags:
rec.tags.append(created_tag)
newly_tagged += 1
imported_or_tagged += newly_tagged
# Record this archive as processed (tag may be None if no media)
arch = ArchiveRecord(
filename=archive_path,
file_size=archive_size,
hash=archive_hash,
artist=artist,
tag=created_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 (content-addressed storage)
# =============================================================================
def build_hashed_dest_path(target_dir: str, original_name: str, content_hash: str) -> str:
"""
Always produce a content-addressed destination:
<target>/<name_without_ext>__<hash[:10]><ext>
If that path already exists, append a numeric suffix.
"""
base, ext = os.path.splitext(original_name)
candidate = os.path.join(target_dir, f"{base}__{content_hash[:10]}{ext}")
if not os.path.exists(candidate):
return candidate
# Rare: collision (e.g., concurrent import) → add counter
i = 2
while True:
candidate = os.path.join(target_dir, f"{base}__{content_hash[:10]}_{i}{ext}")
if not os.path.exists(candidate):
return candidate
i += 1
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).
"""
original_name = os.path.basename(src_path)
print(f"[INFO] Processing: {src_path}")
file_size = os.path.getsize(src_path)
# Compute content hash first (authoritative de-dup + used in filepath)
file_hash = calculate_hash(src_path)
existing = ImageRecord.query.filter_by(hash=file_hash).first()
if existing:
print(f"[SKIP] Duplicate by hash: {original_name}")
return (False, None, existing)
# Optional info: same name+size (but different content) → we still import with hashed name
existing_ns = ImageRecord.query.filter_by(filename=original_name, file_size=file_size).first()
if existing_ns:
print(f"[INFO] Same name+size exists but different hash; storing with hashed name: {original_name}")
# Metadata & pHash similarity
metadata = extract_metadata(src_path)
if not metadata["width"] or not metadata["height"]:
print(f"[SKIP] Missing dimension data for {original_name}")
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] {original_name} is visually similar to an existing larger image.")
return (False, phash, None)
# Copy into /images/<artist>/<base>__<hash[:10]><ext>
target_dir = os.path.join(dest_dir, artist)
os.makedirs(target_dir, exist_ok=True)
dest_path = build_hashed_dest_path(target_dir, original_name, file_hash)
shutil.copy2(src_path, dest_path)
# Thumbnails (mirrored for both images and videos)
try:
if dest_path.lower().endswith((".mp4", ".mov")):
thumb_path = generate_video_thumbnail_mirrored(dest_path)
else:
thumb_path = generate_thumbnail(dest_path)
except Exception as e:
print(f"[WARN] Failed to generate thumbnail for {original_name}: {e}")
thumb_path = None
# Create DB record (keep display name as original; filepath is content-addressed)
record = ImageRecord(
filename=original_name,
filepath=dest_path,
thumb_path=thumb_path,
hash=file_hash,
perceptual_hash=str(phash) if phash else None,
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()
)
# Auto-tag with artist:<name> (treat artist as a tag only)
if artist:
artist_tag_name = f"artist:{artist}"
artist_tag = Tag.query.filter_by(name=artist_tag_name).first()
if not artist_tag:
artist_tag = Tag(name=artist_tag_name, kind="artist")
db.session.add(artist_tag)
db.session.flush()
if artist_tag not in record.tags:
record.tags.append(artist_tag)
db.session.add(record)
if commit:
db.session.commit()
return (True, phash, record)
# =============================================================================
# Helpers: Thumbnails, metadata, hashing, similarity
# =============================================================================
def generate_thumbnail(
image_path,
size=(400, 400),
overwrite=False,
prefer_jpeg=False, # False = auto-pick PNG for alpha images; True = always JPEG
jpeg_bg=(0, 0, 0), # background used when flattening to JPEG
):
"""
Save thumbnails mirrored under /images/thumbs/<...>.
If prefer_jpeg is False (default), transparent images save as PNG, others as JPEG.
If prefer_jpeg is True, all thumbs are JPEG with transparency flattened to `jpeg_bg`.
Returns the absolute path to the thumbnail.
"""
images_root = Path("/images").resolve()
image_path = Path(image_path).resolve()
try:
rel_path = image_path.relative_to(images_root)
except ValueError:
raise ValueError(f"{image_path} is not under {images_root}")
# Base path under /images/thumbs with same subfolders/filename (we may change the suffix)
thumb_base = (images_root / "thumbs" / rel_path).with_suffix("") # drop original suffix for control
# Decide output extension/format
with Image.open(image_path) as img:
img = ImageOps.exif_transpose(img) # honor camera rotation
img.thumbnail(size)
has_alpha = _has_alpha(img)
if prefer_jpeg:
# Force JPEG for consistency
out_path = thumb_base.with_suffix(".jpg")
out_path.parent.mkdir(parents=True, exist_ok=True)
if not overwrite and out_path.exists():
return str(out_path)
img_rgb = _flatten_to_rgb(img, bg=jpeg_bg)
img_rgb.save(out_path, format="JPEG", quality=85, optimize=True, progressive=True)
return str(out_path)
else:
# Auto-pick: PNG for alpha, JPEG otherwise
if has_alpha:
out_path = thumb_base.with_suffix(".png")
out_path.parent.mkdir(parents=True, exist_ok=True)
if not overwrite and out_path.exists():
return str(out_path)
# Keep transparency if present
# Convert palette to RGBA for reliable PNG writing
if img.mode == "P":
img = img.convert("RGBA")
img.save(out_path, format="PNG", optimize=True)
return str(out_path)
else:
out_path = thumb_base.with_suffix(".jpg")
out_path.parent.mkdir(parents=True, exist_ok=True)
if not overwrite and out_path.exists():
return str(out_path)
if img.mode != "RGB":
img = img.convert("RGB")
img.save(out_path, format="JPEG", quality=85, optimize=True, progressive=True)
return str(out_path)
def generate_video_thumbnail(video_path: str, output_path: str, time_position: str = "00:00:01") -> str | None:
"""
Extract a single JPG frame from a video using ffmpeg and write it to output_path.
Returns output_path on success, or None on failure.
"""
try:
os.makedirs(os.path.dirname(output_path), exist_ok=True)
cmd = [
"ffmpeg", "-y",
"-ss", time_position,
"-i", str(video_path),
"-vframes", "1",
"-vf", f"scale={THUMB_SIZE[0]}:-1",
str(output_path),
]
subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=FFMPEG_THUMB_TIMEOUT)
return output_path
except subprocess.CalledProcessError as e:
print(f"[WARN] Failed to extract video thumbnail for {video_path}: {e}")
return None
def generate_video_thumbnail_mirrored(video_path: str, time_position: str = "00:00:01") -> str | None:
"""
Mirror the stored (hashed) video path under /images/thumbs/... and save as .jpg.
"""
images_root = Path("/images").resolve()
video_path_p = Path(video_path).resolve()
rel_path = video_path_p.relative_to(images_root) # raises if not under /images
thumb_path = (images_root / "thumbs" / rel_path).with_suffix(".jpg")
return generate_video_thumbnail(str(video_path_p), str(thumb_path), time_position=time_position)
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
def _has_alpha(img: Image.Image) -> bool:
"""Return True if the image has any transparency channel."""
if img.mode in ("RGBA", "LA"):
return True
if img.mode == "P" and "transparency" in img.info:
return True
return False
def _flatten_to_rgb(img: Image.Image, bg=(0, 0, 0)) -> Image.Image:
"""
Flatten an RGBA/LA/Palette-with-alpha image onto a solid background,
returning an RGB image suitable for saving as JPEG.
"""
# Normalize to RGBA so we can use the alpha channel as a mask
if img.mode == "P":
img = img.convert("RGBA")
if img.mode in ("RGBA", "LA"):
# Ensure we have an explicit alpha channel as last band
if img.mode == "LA":
img = img.convert("RGBA")
alpha = img.split()[-1] # A channel
bg_img = Image.new("RGB", img.size, bg)
bg_img.paste(img, mask=alpha)
return bg_img
# No alpha, just convert to RGB
return img.convert("RGB")
# =============================================================================
# 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
def compute_next_archive_tag_name(artist: str, width: int = ARCHIVE_NUM_WIDTH) -> str:
"""
Return a unique incremental tag name like: 'archive:{artist}/0001', '0002', ...
- Looks only at Tag(kind='archive') with names starting 'archive:{artist}/'
- Finds max numeric suffix and returns next number (zero-padded).
- Ensures no collision even if tags were renamed.
"""
prefix = f"archive:{artist}/"
# Pull existing tag names for this artist/prefix
rows = (Tag.query
.with_entities(Tag.name)
.filter(Tag.kind == 'archive',
Tag.name.like(prefix + '%'))
.all())
max_n = 0
for (name,) in rows:
tail = name[len(prefix):]
if tail.isdigit():
try:
n = int(tail)
if n > max_n:
max_n = n
except Exception:
pass
# Propose next number and make sure it doesn't already exist
n = max_n + 1
while True:
candidate = f"{prefix}{str(n).zfill(width)}"
if not Tag.query.filter_by(name=candidate).first():
return candidate
n += 1
def compute_next_archive_tag_name(artist: str, width: int = ARCHIVE_NUM_WIDTH) -> str:
"""
Return a unique incremental tag name like: 'archive:{artist}/0001', '0002', ...
- Looks only at Tag(kind='archive') with names starting 'archive:{artist}/'
- Finds max numeric suffix and returns next number (zero-padded).
- Ensures no collision even if tags were renamed.
"""
prefix = f"archive:{artist}/"
# Pull existing tag names for this artist/prefix
rows = (Tag.query
.with_entities(Tag.name)
.filter(Tag.kind == 'archive',
Tag.name.like(prefix + '%'))
.all())
max_n = 0
for (name,) in rows:
tail = name[len(prefix):]
if tail.isdigit():
try:
n = int(tail)
if n > max_n:
max_n = n
except Exception:
pass
# Propose next number and make sure it doesn't already exist
n = max_n + 1
while True:
candidate = f"{prefix}{str(n).zfill(width)}"
if not Tag.query.filter_by(name=candidate).first():
return candidate
n += 1