From 0a9ddbf8afb12bfb9862107c43b3be37d5226a46 Mon Sep 17 00:00:00 2001 From: Abdussamet Kocak Date: Thu, 25 Jun 2026 08:55:13 +0300 Subject: [PATCH] feat(download-imageset): add unified download and image postprocess flow --- download_imageset.py | 458 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 458 insertions(+) create mode 100755 download_imageset.py diff --git a/download_imageset.py b/download_imageset.py new file mode 100755 index 0000000..4f8c639 --- /dev/null +++ b/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("> 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(" 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()