feat(download-imageset): add unified download and image postprocess flow

download_imageset.py

@@ -0,0 +1,458 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["httpx", "pillow"]
+# ///
+import argparse
+import dataclasses
+import logging
+import os
+import re
+import struct
+import subprocess
+import sys
+import urllib.parse
+from functools import partial
+from pathlib import Path
+from typing import Callable
+
+import httpx
+from PIL import Image
+
+run_command = partial(
+    subprocess.run,
+    check=True,
+)
+
+# ---------------------------------------------------------------------------
+# Image dimension reading (adapted from sort_images.py)
+# ---------------------------------------------------------------------------
+
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".gif"}
+
+
+def read_dims_jpeg(image_path: Path) -> tuple[int, int]:
+    with image_path.open("rb") as f:
+        if f.read(2) != b"\xff\xd8":
+            raise ValueError(f"{image_path} is not a valid JPEG file")
+        f.seek(0)
+        try:
+            return Image.open(f).size
+        except Exception:
+            f.seek(0)
+            while True:
+                marker = f.read(1)
+                if not marker or marker != b"\xff":
+                    raise ValueError(f"Invalid JPEG format in {image_path}")
+                marker_type = int.from_bytes(f.read(1), byteorder="big")
+                length = int.from_bytes(f.read(2), byteorder="big") - 2
+                is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (
+                    0xC4,
+                    0xC8,
+                    0xCC,
+                )
+                if is_sof:
+                    f.seek(1, 1)
+                    height = int.from_bytes(f.read(2), byteorder="big")
+                    width = int.from_bytes(f.read(2), byteorder="big")
+                    return width, height
+                f.seek(length, 1)
+
+
+def read_dims_png(image_path: Path) -> tuple[int, int]:
+    with image_path.open("rb") as f:
+        if f.read(8) != b"\x89PNG\r\n\x1a\n":
+            raise ValueError(f"{image_path} is not a valid PNG file")
+        try:
+            f.seek(0)
+            with Image.open(f) as img:
+                return img.size
+        except Exception:
+            f.seek(12)
+            if f.read(4) != b"IHDR":
+                raise ValueError(f"IHDR chunk not found in {image_path}")
+            dims = f.read(8)
+            if len(dims) < 8:
+                raise ValueError(f"Truncated IHDR in {image_path}")
+            return struct.unpack(">II", dims)
+
+
+def read_dims_webp(image_path: Path) -> tuple[int, int]:
+    with image_path.open("rb") as f:
+        header = f.read(12)
+        if len(header) < 12 or header[:4] != b"RIFF" or header[8:12] != b"WEBP":
+            raise ValueError(f"Not a valid WebP file: {image_path}")
+        while True:
+            chunk_header = f.read(8)
+            if len(chunk_header) < 8:
+                break
+            tag, length = struct.unpack("<4sI", chunk_header)
+            if tag == b"VP8X":
+                data = f.read(10)
+                return (int.from_bytes(data[4:7], "little") & 0xFFFFFF) + 1, (
+                    int.from_bytes(data[7:10], "little") & 0xFFFFFF
+                ) + 1
+            elif tag == b"VP8L":
+                data = f.read(5)
+                if data[0] != 0x2F:
+                    raise ValueError("Invalid VP8L signature")
+                bits = struct.unpack("<I", data[1:5])[0]
+                return (bits & 0x3FFF) + 1, ((bits >> 14) & 0x3FFF) + 1
+            elif tag == b"VP8 ":
+                data = f.read(10)
+                if data[3:6] != b"\x9d\x01\x2a":
+                    raise ValueError("Invalid VP8 sync code")
+                w, h = struct.unpack("<HH", data[6:10])
+                return w & 0x3FFF, h & 0x3FFF
+            f.seek((length + 1) & ~1, 1)
+    raise ValueError(f"Could not find dimension chunks in {image_path}")
+
+
+def read_dims(image_path: Path) -> tuple[int, int]:
+    match image_path.suffix.lower():
+        case ".jpg" | ".jpeg":
+            return read_dims_jpeg(image_path)
+        case ".png":
+            return read_dims_png(image_path)
+        case ".webp":
+            return read_dims_webp(image_path)
+        case _:
+            raise ValueError(f"Unsupported image format: {image_path.suffix}")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def resolve_download_url_with_alldebrid(token: str, url: str) -> str:
+    resp = httpx.get(
+        "https://api.alldebrid.com/v4/link/unlock",
+        params={"link": url},
+        headers={"Authorization": f"Bearer {token}"},
+        timeout=10,
+    )
+    resp.raise_for_status()
+    data = resp.json()
+    if data["status"] == "error":
+        raise ValueError(f"{data['error']['code']}: {data['error']['message']}")
+    return data["data"]["link"]
+
+
+_CT_TO_EXT = {
+    "application/zip": ".zip",
+    "application/x-zip-compressed": ".zip",
+    "application/x-rar-compressed": ".rar",
+    "application/vnd.rar": ".rar",
+    "application/x-7z-compressed": ".7z",
+    "video/mp4": ".mp4",
+    "video/quicktime": ".mov",
+    "video/x-matroska": ".mkv",
+}
+
+_VIDEO_EXTS = {".mp4", ".mov", ".mkv"}
+
+
+@dataclasses.dataclass
+class UrlMeta:
+    suggested_filename: str
+    suggested_extension: str
+    size_bytes: int
+
+
+def _resolve_url(url: str) -> UrlMeta:
+    """Send HEAD request; return filename, extension, and size from response headers."""
+    try:
+        resp = httpx.head(url, follow_redirects=True, timeout=15)
+        cd = resp.headers.get("content-disposition", "")
+        filename = None
+        if "filename=" in cd:
+            part = cd.split("filename=", 1)[1].strip().strip('"').strip("'")
+            filename = urllib.parse.unquote(part.split(";")[0].strip())
+
+        ct = resp.headers.get("content-type", "").split(";")[0].strip()
+        ext = _CT_TO_EXT.get(ct) or Path(urllib.parse.urlparse(url).path).suffix
+
+        if not filename:
+            filename = urllib.parse.unquote(Path(urllib.parse.urlparse(url).path).name)
+        if not Path(filename).suffix and ext:
+            filename = filename + ext
+
+        size = int(resp.headers.get("content-length", 0))
+        return UrlMeta(
+            suggested_filename=filename, suggested_extension=ext, size_bytes=size
+        )
+    except Exception:
+        url_path = Path(urllib.parse.urlparse(url).path)
+        name = urllib.parse.unquote(url_path.name)
+        return UrlMeta(
+            suggested_filename=name, suggested_extension=url_path.suffix, size_bytes=0
+        )
+
+
+def _name_from_url(url: str) -> str:
+    parsed = urllib.parse.urlparse(url)
+    segment = parsed.path.rstrip("/").split("/")[-1]
+    return urllib.parse.unquote(segment) or parsed.netloc
+
+
+def _has_images(d: Path) -> bool:
+    return any(f.suffix.lower() in IMAGE_EXTS for f in d.rglob("*") if f.is_file())
+
+
+def _needs_resize(d: Path) -> bool:
+    images = [f for f in d.rglob("*") if f.is_file() and f.suffix.lower() in IMAGE_EXTS]
+    if not images:
+        return False
+    mean_mb = sum(f.stat().st_size for f in images) / len(images)
+    files_are_large = mean_mb >= 2 * 1_048_567
+    images_are_large = any(max(read_dims(img)) > 5000 for img in images)
+
+    return files_are_large or images_are_large
+
+
+# ---------------------------------------------------------------------------
+# Post-download: extract archives then resize
+# ---------------------------------------------------------------------------
+
+
+def _extract_and_process(output_dir: Path) -> None:
+    archives = [
+        f for pattern in ("*.zip", "*.rar", "*.7z") for f in output_dir.glob(pattern)
+    ]
+    for archive in archives:
+        print(f"extracting {archive.name} …")
+        run_command(["unar", "-o", str(output_dir), str(archive)])
+        move_to_trash([archive])
+
+    subdirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name != "resized"]
+    targets = subdirs if subdirs else [output_dir]
+
+    to_process = [d for d in targets if _needs_resize(d)]
+    if not to_process:
+        print("no dirs need resizing")
+    else:
+        resized_dir = output_dir / "resized"
+        resized_dir.mkdir(exist_ok=True)
+        originals = [
+            f
+            for d in to_process
+            for f in d.rglob("*")
+            if f.is_file() and f.suffix.lower() in IMAGE_EXTS
+        ]
+        print(f"resizing {len(originals)} image(s) → {resized_dir}")
+        for img in originals:
+            run_command(
+                [
+                    "imgz",
+                    "resize",
+                    "-o",
+                    str(resized_dir / f"{img.stem}.resized.jpg"),
+                    str(img),
+                ]
+            )
+        if originals:
+            print(f"trashing {len(originals)} original(s)")
+            move_to_trash(originals)
+    rename_children(output_dir, flatten=True)
+    remove_empty_dirs(output_dir)
+
+
+def remove_empty_dirs(root_dir: Path) -> None:
+    removed = [
+        d
+        for d in sorted(root_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True)
+        if d.is_dir() and not any(d.iterdir())
+    ]
+    for d in removed:
+        d.rmdir()
+    if removed:
+        print(f"removed {len(removed)} empty dir(s) in {root_dir}")
+
+
+def move_to_trash(file_paths: list[Path]) -> None:
+    """Move files to Trash using trash CLI (single batched call)."""
+    resolved: list[str] = []
+    for file_path in file_paths:
+        if not file_path.is_file():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        resolved.append(str(file_path.resolve()))
+
+    if resolved:
+        run_command(["trash", *resolved])
+
+
+def rename_children(
+    source_dir: Path,
+    flatten: bool,
+    filter_fn: Callable[[Path], bool] | None = None,
+) -> None:
+    if not filter_fn:
+        filter_fn = is_image
+
+    def natural_sort_key(path: Path) -> tuple[str, ...]:
+        s = path.stem
+        return tuple(
+            f"{int(p):010d}" if p.isnumeric() else p
+            for p in re.findall(r"(\D+|\d+)", s)
+        )
+
+    assert source_dir.is_dir()
+
+    if flatten:
+        files = [f for f in source_dir.rglob("*") if f.is_file()]
+    else:
+        files = [f for f in source_dir.iterdir() if f.is_file()]
+
+    files = [f for f in files if filter_fn(f)]
+
+    files = sorted(files, key=natural_sort_key)
+
+    for i, f in enumerate(files):
+        ext = ".jpg" if f.suffix.lower() in {".jpeg", ".jpg"} else f.suffix
+        target_dir = f.parent
+        if flatten:
+            target_dir = source_dir
+        new_stem = source_dir.name
+        target = target_dir / f"{new_stem}__{i:04d}{ext}"
+        if target.is_file():
+            logging.error(f"target already exists: {target}")
+            continue
+        f.rename(target)
+
+
+def is_image(f: Path) -> bool:
+    return f.is_file() and f.suffix.lower() in IMAGE_EXTS
+
+
+# ---------------------------------------------------------------------------
+# Subcommands
+# ---------------------------------------------------------------------------
+
+
+def cmd_gallery(urls: list[str], name: str | None, cwd: Path | None = None) -> None:
+    resolved_name = name or _name_from_url(urls[0])
+    output_dir = (cwd or Path(".")) / resolved_name
+
+    if output_dir.exists() and _has_images(output_dir):
+        print(f"skipping gallery-dl: {output_dir} already has images")
+    else:
+        print(f"gallery-dl → {output_dir}")
+        args = ["uvx", "gallery-dl"]
+        if name or cwd:
+            args += ["-D", str(output_dir)]
+        run_command(args + urls)
+
+    _extract_and_process(output_dir)
+
+
+def cmd_filehost(
+    urls: list[str], name: str | None, token: str, cwd: Path | None = None
+) -> None:
+    resolved_name = name or _name_from_url(urls[0])
+    output_dir = (cwd or Path(".")) / resolved_name
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for url in urls:
+        print(f"unrestricting {url}")
+        dl_url = resolve_download_url_with_alldebrid(token, url)
+        meta = _resolve_url(dl_url)
+        if meta.suggested_extension in _VIDEO_EXTS:
+            video_path = (cwd or Path(".")) / (resolved_name + meta.suggested_extension)
+            print(f"video detected → {video_path}")
+            if video_path.exists() and video_path.stat().st_size > 5 * 1024 * 1024:
+                print(f"skipping: {video_path.name} already downloaded")
+                continue
+            run_command(["aria2c", "-o", str(video_path), dl_url])
+            continue
+
+        dest = output_dir / meta.suggested_filename
+        if dest.exists() and dest.stat().st_size > 5 * 1024 * 1024:
+            print(f"skipping: {dest.name} already downloaded")
+            continue
+        print(f"downloading {meta.suggested_filename} → {output_dir}")
+        run_command(["aria2c", f"--dir={output_dir}", dl_url])
+
+    _extract_and_process(output_dir)
+
+
+def cmd_process(dirs: list[str]) -> None:
+    for d in dirs:
+        p = Path(d)
+        if not p.exists():
+            print(f"skipping {p}: does not exist", file=sys.stderr)
+            continue
+        if _needs_resize(p):
+            resized_dir = p / "resized"
+            resized_dir.mkdir(exist_ok=True)
+            originals = [f for f in p.rglob("*") if is_image(f)]
+            print(f"resizing {len(originals)} image(s) in {p} → {resized_dir}")
+            for img in originals:
+                run_command(
+                    [
+                        "imgz",
+                        "resize",
+                        "-o",
+                        str(resized_dir / f"{img.stem}.resized.jpg"),
+                        str(img),
+                    ]
+                )
+            if originals:
+                print(f"trashing {len(originals)} original(s)")
+                move_to_trash(originals)
+            remove_empty_dirs(p)
+        else:
+            print(f"skipping {p}: already within size limits")
+
+
+def cmd_download(
+    urls: list[str], name: str | None, token: str, cwd: Path | None = None
+) -> None:
+    filehost_urls = [u for u in urls if "rg.to" in u or "rapidgator.net" in u]
+    gallery_urls = [u for u in urls if u not in set(filehost_urls)]
+    if filehost_urls:
+        cmd_filehost(filehost_urls, name, token, cwd)
+    if gallery_urls:
+        cmd_gallery(gallery_urls, name, cwd)
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    print(os.getenv("PATH"))
+    parser = argparse.ArgumentParser(prog="download_imageset.py")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    gallery_p = sub.add_parser("gallery")
+    gallery_p.add_argument("urls", nargs="+")
+    gallery_p.add_argument("--name", default=None)
+    gallery_p.add_argument("--cwd", type=Path, default=None)
+
+    for cmd in ("filehost", "download"):
+        p = sub.add_parser(cmd)
+        p.add_argument("urls", nargs="+")
+        p.add_argument("--name", default=None)
+        p.add_argument("--cwd", default=None, type=Path)
+        p.add_argument("--token", default=os.getenv("ALLDEBRID_TOKEN"))
+
+    proc = sub.add_parser("process")
+    proc.add_argument("dirs", nargs="+")
+
+    args = parser.parse_args()
+
+    match args.command:
+        case "gallery":
+            cmd_gallery(args.urls, args.name, args.cwd)
+        case "filehost":
+            cmd_filehost(args.urls, args.name, args.token, args.cwd)
+        case "process":
+            cmd_process(args.dirs)
+        case "download":
+            cmd_download(args.urls, args.name, args.token, args.cwd)
+
+
+if __name__ == "__main__":
+    main()