349 lines
11 KiB
Python
349 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import io
|
||
import re
|
||
import sys
|
||
import zipfile
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from dataclasses import dataclass
|
||
from pathlib import Path, PurePosixPath
|
||
from typing import Callable
|
||
|
||
try:
|
||
from PIL import Image, ImageSequence
|
||
except ImportError as exc: # pragma: no cover
|
||
raise SystemExit(
|
||
"This script requires Pillow. Install it with: python3 -m pip install pillow"
|
||
) from exc
|
||
|
||
from animated_webp_to_gif import webp_bytes_to_gif
|
||
|
||
IMAGE_EXTENSIONS = {
|
||
".avif",
|
||
".bmp",
|
||
".gif",
|
||
".jfif",
|
||
".jpeg",
|
||
".jpg",
|
||
".jpe",
|
||
".png",
|
||
".tif",
|
||
".tiff",
|
||
".webp",
|
||
}
|
||
|
||
# 已编号的成品:1.webp、42.gif(不参与转换;用于各子文件夹内续号)
|
||
OUTPUT_NAME_RE = re.compile(r"^(\d+)\.(webp|gif)$", re.IGNORECASE)
|
||
|
||
DEFAULT_SOURCE_DIR = Path("原图")
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Job:
|
||
source_label: str
|
||
target_dir: Path
|
||
index: int
|
||
loader: Callable[[], bytes]
|
||
# 磁盘上的源文件;转换成功后删除(zip 成员为 None)。与输出路径相同时不会删。
|
||
source_file: Path | None = None
|
||
|
||
|
||
def is_image_path(path: Path) -> bool:
|
||
return path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
|
||
|
||
|
||
def is_numbered_asset_file(path: Path) -> bool:
|
||
return path.is_file() and bool(OUTPUT_NAME_RE.match(path.name))
|
||
|
||
|
||
def max_index_in_directory(directory: Path) -> int:
|
||
highest = 0
|
||
if not directory.is_dir():
|
||
return 0
|
||
for p in directory.iterdir():
|
||
if not p.is_file():
|
||
continue
|
||
m = OUTPUT_NAME_RE.match(p.name)
|
||
if m:
|
||
highest = max(highest, int(m.group(1)))
|
||
return highest
|
||
|
||
|
||
def convert_to_webp(data: bytes) -> tuple[bytes, bool]:
|
||
with Image.open(io.BytesIO(data)) as im:
|
||
is_animated = bool(
|
||
getattr(im, "is_animated", False) or getattr(im, "n_frames", 1) > 1
|
||
)
|
||
|
||
if is_animated:
|
||
frames = []
|
||
durations = []
|
||
loop = im.info.get("loop", 0)
|
||
for frame in ImageSequence.Iterator(im):
|
||
frames.append(frame.convert("RGBA"))
|
||
durations.append(frame.info.get("duration", im.info.get("duration", 100)))
|
||
|
||
output = io.BytesIO()
|
||
frames[0].save(
|
||
output,
|
||
format="WEBP",
|
||
save_all=True,
|
||
append_images=frames[1:],
|
||
duration=durations,
|
||
loop=loop,
|
||
lossless=True,
|
||
method=4,
|
||
)
|
||
return output.getvalue(), True
|
||
|
||
output = io.BytesIO()
|
||
im.convert("RGBA").save(output, format="WEBP", lossless=True, method=4)
|
||
return output.getvalue(), False
|
||
|
||
|
||
def is_animated_gif(data: bytes) -> bool:
|
||
with Image.open(io.BytesIO(data)) as im:
|
||
if im.format != "GIF":
|
||
return False
|
||
n_frames = getattr(im, "n_frames", 1)
|
||
return bool(getattr(im, "is_animated", False) or n_frames > 1)
|
||
|
||
|
||
def is_animated_webp_bytes(data: bytes) -> bool:
|
||
with Image.open(io.BytesIO(data)) as im:
|
||
if im.format != "WEBP":
|
||
return False
|
||
return bool(
|
||
getattr(im, "is_animated", False) or getattr(im, "n_frames", 1) > 1
|
||
)
|
||
|
||
|
||
def _read_zip_member(zip_path: Path, member_name: str) -> bytes:
|
||
with zipfile.ZipFile(zip_path) as zf:
|
||
return zf.read(member_name)
|
||
|
||
|
||
def collect_directory_jobs(
|
||
source_root: Path, output_root: Path, *, delete_sources: bool
|
||
) -> list[Job]:
|
||
grouped: dict[Path, list[Path]] = {}
|
||
for path in sorted(source_root.rglob("*")):
|
||
if not is_image_path(path):
|
||
continue
|
||
if is_numbered_asset_file(path):
|
||
continue
|
||
grouped.setdefault(path.parent, []).append(path)
|
||
|
||
jobs: list[Job] = []
|
||
for directory in sorted(grouped, key=lambda p: str(p)):
|
||
files = sorted(grouped[directory], key=lambda p: p.name)
|
||
target_dir = output_root / directory.relative_to(source_root)
|
||
next_i = max_index_in_directory(target_dir) + 1
|
||
for path in files:
|
||
jobs.append(
|
||
Job(
|
||
source_label=path.relative_to(source_root).as_posix(),
|
||
target_dir=target_dir,
|
||
index=next_i,
|
||
loader=lambda p=path: p.read_bytes(),
|
||
source_file=path if delete_sources else None,
|
||
)
|
||
)
|
||
next_i += 1
|
||
return jobs
|
||
|
||
|
||
def collect_zip_jobs(zip_path: Path, output_root: Path) -> list[Job]:
|
||
grouped: dict[PurePosixPath, list[PurePosixPath]] = {}
|
||
with zipfile.ZipFile(zip_path) as zf:
|
||
for info in zf.infolist():
|
||
if info.is_dir():
|
||
continue
|
||
member_path = PurePosixPath(info.filename)
|
||
suf = member_path.suffix.lower()
|
||
if suf not in IMAGE_EXTENSIONS:
|
||
continue
|
||
if OUTPUT_NAME_RE.match(member_path.name):
|
||
continue
|
||
grouped.setdefault(member_path.parent, []).append(member_path)
|
||
|
||
jobs: list[Job] = []
|
||
for directory in sorted(grouped, key=lambda p: str(p)):
|
||
files = sorted(grouped[directory], key=lambda p: str(p.name))
|
||
target_dir = output_root / Path(str(directory))
|
||
next_i = max_index_in_directory(target_dir) + 1
|
||
for member_path in files:
|
||
jobs.append(
|
||
Job(
|
||
source_label=member_path.as_posix(),
|
||
target_dir=target_dir,
|
||
index=next_i,
|
||
loader=lambda name=member_path.as_posix(), zp=zip_path: _read_zip_member(
|
||
zp, name
|
||
),
|
||
)
|
||
)
|
||
next_i += 1
|
||
return jobs
|
||
|
||
|
||
def _remove_source_if_any(job: Job, out: Path, kind: str) -> str:
|
||
if job.source_file is None:
|
||
return kind
|
||
try:
|
||
src = job.source_file.resolve()
|
||
if src == out.resolve():
|
||
return kind
|
||
src.unlink()
|
||
return f"{kind}; removed-source"
|
||
except OSError as exc:
|
||
print(
|
||
f"[warn] ok: {out} ({kind}) but could not delete {job.source_file}: {exc}",
|
||
file=sys.stderr,
|
||
)
|
||
return kind
|
||
|
||
|
||
def run_job(job: Job) -> tuple[bool, str, Path, str | None, str | None]:
|
||
try:
|
||
data = job.loader()
|
||
job.target_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
if is_animated_gif(data):
|
||
out = job.target_dir / f"{job.index}.gif"
|
||
out.write_bytes(data)
|
||
kind = _remove_source_if_any(job, out, "gif-animated")
|
||
return True, job.source_label, out, kind, None
|
||
|
||
if is_animated_webp_bytes(data):
|
||
out = job.target_dir / f"{job.index}.gif"
|
||
webp_bytes_to_gif(data, out)
|
||
kind = _remove_source_if_any(job, out, "animated-webp->gif")
|
||
return True, job.source_label, out, kind, None
|
||
|
||
webp_data, animated = convert_to_webp(data)
|
||
out = job.target_dir / f"{job.index}.webp"
|
||
out.write_bytes(webp_data)
|
||
base = "animated-webp" if animated else "static-webp"
|
||
kind = _remove_source_if_any(job, out, base)
|
||
return True, job.source_label, out, kind, None
|
||
except Exception as exc:
|
||
return False, job.source_label, job.target_dir / f"{job.index}.?", None, str(exc)
|
||
|
||
|
||
def main(argv: list[str]) -> int:
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"Read unnumbered images from 原图/ (by default), write numbered *.webp / *.gif "
|
||
"into the same folder tree (in-place). Subfolders mirror 原图. "
|
||
"Animated GIFs and animated WebPs become .gif; static images become lossless .webp. "
|
||
"Next index per folder = max existing N in that folder + 1. "
|
||
"By default unnumbered source files are deleted after a successful conversion."
|
||
)
|
||
)
|
||
parser.add_argument(
|
||
"source",
|
||
nargs="?",
|
||
default=None,
|
||
type=Path,
|
||
help=f"Source root folder, single image, or .zip (default: {DEFAULT_SOURCE_DIR.as_posix()}/).",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--output",
|
||
default=None,
|
||
dest="output_root",
|
||
type=Path,
|
||
help="Output root (default: same as source for folders/single file; for .zip default: <zip-stem>/).",
|
||
)
|
||
parser.add_argument(
|
||
"-j",
|
||
"--jobs",
|
||
type=int,
|
||
default=10,
|
||
help="Number of worker threads (default: 10).",
|
||
)
|
||
parser.add_argument(
|
||
"--keep-sources",
|
||
action="store_true",
|
||
help="Keep unnumbered source files after conversion (default: delete them).",
|
||
)
|
||
args = parser.parse_args(argv)
|
||
|
||
cwd = Path.cwd()
|
||
source = (
|
||
args.source.expanduser().resolve()
|
||
if args.source is not None
|
||
else (cwd / DEFAULT_SOURCE_DIR).resolve()
|
||
)
|
||
jobs: list[Job] = []
|
||
|
||
if source.is_dir():
|
||
out_root = (
|
||
args.output_root.expanduser().resolve()
|
||
if args.output_root is not None
|
||
else source
|
||
)
|
||
out_root.mkdir(parents=True, exist_ok=True)
|
||
jobs = collect_directory_jobs(
|
||
source, out_root, delete_sources=not args.keep_sources
|
||
)
|
||
elif source.is_file() and source.suffix.lower() == ".zip":
|
||
out_root = (
|
||
args.output_root.expanduser().resolve()
|
||
if args.output_root is not None
|
||
else (source.parent / source.stem).resolve()
|
||
)
|
||
out_root.mkdir(parents=True, exist_ok=True)
|
||
jobs = collect_zip_jobs(source, out_root)
|
||
elif is_image_path(source):
|
||
if is_numbered_asset_file(source):
|
||
print(f"Skip already-numbered file: {source}", file=sys.stderr)
|
||
return 0
|
||
target_dir = (
|
||
args.output_root.expanduser().resolve()
|
||
if args.output_root is not None
|
||
else source.parent.resolve()
|
||
)
|
||
target_dir.mkdir(parents=True, exist_ok=True)
|
||
idx = max_index_in_directory(target_dir) + 1
|
||
jobs = [
|
||
Job(
|
||
source_label=source.name,
|
||
target_dir=target_dir,
|
||
index=idx,
|
||
loader=lambda p=source: p.read_bytes(),
|
||
source_file=source if not args.keep_sources else None,
|
||
)
|
||
]
|
||
else:
|
||
print(f"Not a folder, image, or zip: {source}", file=sys.stderr)
|
||
return 1
|
||
|
||
jobs.sort(key=lambda job: (str(job.target_dir), job.index))
|
||
|
||
if not jobs:
|
||
print(f"No new images to convert under: {source}")
|
||
return 0
|
||
|
||
converted = 0
|
||
skipped = 0
|
||
|
||
with ThreadPoolExecutor(max_workers=max(1, args.jobs)) as pool:
|
||
for ok, label, target, kind, error in pool.map(run_job, jobs):
|
||
if ok:
|
||
print(f"{label} -> {target} ({kind})")
|
||
converted += 1
|
||
else:
|
||
print(f"[skip] {label}: {error}", file=sys.stderr)
|
||
skipped += 1
|
||
|
||
print(f"Done. converted={converted}, skipped={skipped}")
|
||
return 0 if converted else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main(sys.argv[1:]))
|