Files
meme-api/convert_to_webp.py

349 lines
11 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.
#!/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:]))