feat(download-imageset): add unified download and image postprocess flow
This commit is contained in:
Executable
+458
@@ -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()
|
||||
Reference in New Issue
Block a user