feat: optimize compression and production deployment
This commit is contained in:
237899745
404
scripts/benchmark_compression.py
Normal file
404
scripts/benchmark_compression.py
Normal file
@@ -0,0 +1,404 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate a real-photo corpus and benchmark ImageForge through its HTTP API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from http.cookiejar import CookieJar
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
|
||||
CORPUS = (
|
||||
("photo_1015_jpeg", 1015, "JPEG", "jpg"),
|
||||
("photo_1016_png", 1016, "PNG", "png"),
|
||||
("photo_1025_webp", 1025, "WEBP", "webp"),
|
||||
("photo_1039_avif", 1039, "AVIF", "avif"),
|
||||
)
|
||||
FORMAT_EXTENSIONS = {"jpeg": "jpg", "webp": "webp", "avif": "avif"}
|
||||
PIL_FORMATS = {"jpeg": "JPEG", "webp": "WEBP", "avif": "AVIF"}
|
||||
|
||||
|
||||
def parse_csv_arg(value: str, cast: type = str) -> list[Any]:
|
||||
return [cast(item.strip()) for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def open_image(path: Path) -> Image.Image:
|
||||
with Image.open(path) as image:
|
||||
return ImageOps.exif_transpose(image).convert("RGB")
|
||||
|
||||
|
||||
def download(url: str) -> bytes:
|
||||
request = urllib.request.Request(url, headers={"User-Agent": "ImageForge benchmark/1.0"})
|
||||
with urllib.request.urlopen(request, timeout=90) as response:
|
||||
return response.read()
|
||||
|
||||
|
||||
def generate_corpus(output_dir: Path, width: int, height: int) -> None:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
manifest: list[dict[str, Any]] = []
|
||||
|
||||
for name, image_id, image_format, extension in CORPUS:
|
||||
source_url = f"https://picsum.photos/id/{image_id}/{width}/{height}.jpg"
|
||||
source_bytes = download(source_url)
|
||||
source_path = output_dir / f".{name}.source.jpg"
|
||||
source_path.write_bytes(source_bytes)
|
||||
image = open_image(source_path)
|
||||
source_path.unlink()
|
||||
|
||||
# Normalize dimensions so format and content complexity, not resolution, drive comparisons.
|
||||
image = ImageOps.fit(image, (width, height), method=Image.Resampling.LANCZOS)
|
||||
output_path = output_dir / f"{name}.{extension}"
|
||||
if image_format == "JPEG":
|
||||
image.save(output_path, format=image_format, quality=95, subsampling=0, optimize=True)
|
||||
elif image_format == "PNG":
|
||||
image.save(output_path, format=image_format, optimize=True, compress_level=9)
|
||||
elif image_format == "WEBP":
|
||||
image.save(output_path, format=image_format, lossless=True, method=6)
|
||||
else:
|
||||
image.save(output_path, format=image_format, quality=95, speed=6)
|
||||
|
||||
manifest.append(
|
||||
{
|
||||
"file": output_path.name,
|
||||
"source_url": source_url,
|
||||
"picsum_id": image_id,
|
||||
"format": image_format.lower(),
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"size_bytes": output_path.stat().st_size,
|
||||
}
|
||||
)
|
||||
|
||||
(output_dir / "manifest.json").write_text(
|
||||
json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
print(f"generated {len(manifest)} images in {output_dir}")
|
||||
|
||||
|
||||
def multipart_body(file_path: Path, fields: dict[str, str]) -> tuple[bytes, str]:
|
||||
boundary = f"----imageforge-{uuid.uuid4().hex}"
|
||||
chunks: list[bytes] = []
|
||||
for name, value in fields.items():
|
||||
chunks.extend(
|
||||
(
|
||||
f"--{boundary}\r\n".encode(),
|
||||
f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
|
||||
value.encode(),
|
||||
b"\r\n",
|
||||
)
|
||||
)
|
||||
|
||||
chunks.extend(
|
||||
(
|
||||
f"--{boundary}\r\n".encode(),
|
||||
(
|
||||
f'Content-Disposition: form-data; name="file"; '
|
||||
f'filename="{file_path.name}"\r\n'
|
||||
).encode(),
|
||||
b"Content-Type: application/octet-stream\r\n\r\n",
|
||||
file_path.read_bytes(),
|
||||
b"\r\n",
|
||||
f"--{boundary}--\r\n".encode(),
|
||||
)
|
||||
)
|
||||
return b"".join(chunks), f"multipart/form-data; boundary={boundary}"
|
||||
|
||||
|
||||
def request_json(
|
||||
opener: urllib.request.OpenerDirector,
|
||||
url: str,
|
||||
data: bytes,
|
||||
headers: dict[str, str],
|
||||
timeout: int,
|
||||
) -> dict[str, Any]:
|
||||
request = urllib.request.Request(url, data=data, headers=headers, method="POST")
|
||||
try:
|
||||
with opener.open(request, timeout=timeout) as response:
|
||||
payload = json.load(response)
|
||||
except urllib.error.HTTPError as error:
|
||||
detail = error.read().decode("utf-8", errors="replace")
|
||||
raise RuntimeError(f"HTTP {error.code}: {detail}") from error
|
||||
|
||||
if not payload.get("success") or "data" not in payload:
|
||||
raise RuntimeError(f"unexpected API response: {payload}")
|
||||
return payload["data"]
|
||||
|
||||
|
||||
def login(
|
||||
opener: urllib.request.OpenerDirector,
|
||||
base_url: str,
|
||||
email: str,
|
||||
password: str,
|
||||
timeout: int,
|
||||
) -> str:
|
||||
body = json.dumps({"email": email, "password": password}).encode()
|
||||
data = request_json(
|
||||
opener,
|
||||
f"{base_url}/api/v1/auth/login",
|
||||
body,
|
||||
{"Content-Type": "application/json"},
|
||||
timeout,
|
||||
)
|
||||
return str(data["token"])
|
||||
|
||||
|
||||
def block_ssim(reference: np.ndarray, candidate: np.ndarray, block: int = 8) -> float:
|
||||
ref = 0.2126 * reference[..., 0] + 0.7152 * reference[..., 1] + 0.0722 * reference[..., 2]
|
||||
out = 0.2126 * candidate[..., 0] + 0.7152 * candidate[..., 1] + 0.0722 * candidate[..., 2]
|
||||
height = (ref.shape[0] // block) * block
|
||||
width = (ref.shape[1] // block) * block
|
||||
if height == 0 or width == 0:
|
||||
height, width, block = ref.shape[0], ref.shape[1], 1
|
||||
|
||||
def blocks(array: np.ndarray) -> np.ndarray:
|
||||
return (
|
||||
array[:height, :width]
|
||||
.reshape(height // block, block, width // block, block)
|
||||
.transpose(0, 2, 1, 3)
|
||||
)
|
||||
|
||||
ref_blocks = blocks(ref)
|
||||
out_blocks = blocks(out)
|
||||
axes = (-1, -2)
|
||||
ref_mean = ref_blocks.mean(axis=axes)
|
||||
out_mean = out_blocks.mean(axis=axes)
|
||||
ref_var = ref_blocks.var(axis=axes)
|
||||
out_var = out_blocks.var(axis=axes)
|
||||
covariance = ((ref_blocks - ref_mean[..., None, None]) * (out_blocks - out_mean[..., None, None])).mean(axis=axes)
|
||||
c1 = (0.01 * 255.0) ** 2
|
||||
c2 = (0.03 * 255.0) ** 2
|
||||
numerator = (2 * ref_mean * out_mean + c1) * (2 * covariance + c2)
|
||||
denominator = (ref_mean**2 + out_mean**2 + c1) * (ref_var + out_var + c2)
|
||||
return float(np.mean(numerator / np.maximum(denominator, 1e-12)))
|
||||
|
||||
|
||||
def image_metrics(reference_path: Path, output_path: Path) -> dict[str, Any]:
|
||||
reference = open_image(reference_path)
|
||||
with Image.open(output_path) as opened:
|
||||
detected_format = (opened.format or "unknown").upper()
|
||||
output = ImageOps.exif_transpose(opened).convert("RGB")
|
||||
|
||||
output_width, output_height = output.size
|
||||
if output.size != reference.size:
|
||||
output = output.resize(reference.size, Image.Resampling.LANCZOS)
|
||||
|
||||
ref_array = np.asarray(reference, dtype=np.float64)
|
||||
out_array = np.asarray(output, dtype=np.float64)
|
||||
mse = float(np.mean((ref_array - out_array) ** 2))
|
||||
psnr = 99.0 if mse == 0 else 20.0 * math.log10(255.0 / math.sqrt(mse))
|
||||
return {
|
||||
"detected_format": detected_format,
|
||||
"width": output_width,
|
||||
"height": output_height,
|
||||
"pixel_ratio_pct": output_width * output_height * 100.0 / (reference.width * reference.height),
|
||||
"ssim": block_ssim(ref_array, out_array),
|
||||
"psnr_db": psnr,
|
||||
}
|
||||
|
||||
|
||||
def summarize(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
groups: dict[tuple[str, int], list[dict[str, Any]]] = {}
|
||||
for row in rows:
|
||||
if row.get("error"):
|
||||
continue
|
||||
groups.setdefault((str(row["output_format"]), int(row["requested_rate"])), []).append(row)
|
||||
|
||||
summary: list[dict[str, Any]] = []
|
||||
for (output_format, requested_rate), group in sorted(groups.items()):
|
||||
summary.append(
|
||||
{
|
||||
"output_format": output_format,
|
||||
"requested_rate": requested_rate,
|
||||
"cases": len(group),
|
||||
"target_met": sum(bool(row["target_met"]) for row in group),
|
||||
"format_ok": sum(bool(row["format_ok"]) for row in group),
|
||||
"mean_actual_rate_pct": statistics.mean(float(row["actual_rate_pct"]) for row in group),
|
||||
"mean_saved_pct": statistics.mean(float(row["saved_pct"]) for row in group),
|
||||
"mean_ssim": statistics.mean(float(row["ssim"]) for row in group),
|
||||
"mean_psnr_db": statistics.mean(float(row["psnr_db"]) for row in group),
|
||||
"mean_pixel_ratio_pct": statistics.mean(float(row["pixel_ratio_pct"]) for row in group),
|
||||
"median_elapsed_ms": statistics.median(float(row["elapsed_ms"]) for row in group),
|
||||
}
|
||||
)
|
||||
return summary
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> int:
|
||||
input_dir = Path(args.input_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
inputs = sorted(path for path in input_dir.iterdir() if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".avif"})
|
||||
if not inputs:
|
||||
raise RuntimeError(f"no benchmark images found in {input_dir}")
|
||||
|
||||
rates = parse_csv_arg(args.rates, int)
|
||||
formats = [str(value).lower() for value in parse_csv_arg(args.formats)]
|
||||
unsupported = sorted(set(formats) - set(FORMAT_EXTENSIONS))
|
||||
if unsupported:
|
||||
raise RuntimeError(f"unsupported output formats: {', '.join(unsupported)}")
|
||||
|
||||
base_url = args.base_url.rstrip("/")
|
||||
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(CookieJar()))
|
||||
token = args.token or os.getenv("IMAGEFORGE_BENCH_TOKEN", "")
|
||||
if not token:
|
||||
email = args.email or os.getenv("IMAGEFORGE_BENCH_EMAIL", "")
|
||||
password = args.password or os.getenv("IMAGEFORGE_BENCH_PASSWORD", "")
|
||||
if email and password:
|
||||
token = login(opener, base_url, email, password, args.timeout)
|
||||
|
||||
auth_headers = {"Authorization": f"Bearer {token}"} if token else {}
|
||||
rows: list[dict[str, Any]] = []
|
||||
total = len(inputs) * len(formats) * len(rates)
|
||||
case_number = 0
|
||||
|
||||
for input_path in inputs:
|
||||
original_size = input_path.stat().st_size
|
||||
for output_format in formats:
|
||||
for rate in rates:
|
||||
case_number += 1
|
||||
output_path = output_dir / f"{input_path.stem}__{output_format}__r{rate}.{FORMAT_EXTENSIONS[output_format]}"
|
||||
row: dict[str, Any] = {
|
||||
"input": input_path.name,
|
||||
"input_format": input_path.suffix.lower().lstrip("."),
|
||||
"output_format": output_format,
|
||||
"requested_rate": rate,
|
||||
"original_size": original_size,
|
||||
}
|
||||
try:
|
||||
body, content_type = multipart_body(
|
||||
input_path,
|
||||
{"compression_rate": str(rate), "output_format": output_format},
|
||||
)
|
||||
started = time.perf_counter()
|
||||
data = request_json(
|
||||
opener,
|
||||
f"{base_url}/api/v1/compress",
|
||||
body,
|
||||
{**auth_headers, "Content-Type": content_type},
|
||||
args.timeout,
|
||||
)
|
||||
elapsed_ms = (time.perf_counter() - started) * 1000.0
|
||||
|
||||
download_request = urllib.request.Request(
|
||||
urllib.parse.urljoin(f"{base_url}/", str(data["download_url"]).lstrip("/")),
|
||||
headers=auth_headers,
|
||||
)
|
||||
with opener.open(download_request, timeout=args.timeout) as response:
|
||||
output_path.write_bytes(response.read())
|
||||
|
||||
compressed_size = output_path.stat().st_size
|
||||
metrics = image_metrics(input_path, output_path)
|
||||
tolerance_bytes = max(1024, int(original_size * 0.01))
|
||||
target_bytes = original_size * rate / 100.0
|
||||
row.update(
|
||||
{
|
||||
"compressed_size": compressed_size,
|
||||
"actual_rate_pct": compressed_size * 100.0 / original_size,
|
||||
"saved_pct": max(0.0, (original_size - compressed_size) * 100.0 / original_size),
|
||||
"target_error_pct_points": compressed_size * 100.0 / original_size - rate,
|
||||
"target_met": compressed_size <= target_bytes + tolerance_bytes,
|
||||
"format_ok": metrics["detected_format"] == PIL_FORMATS[output_format],
|
||||
"api_size_matches": int(data["compressed_size"]) == compressed_size,
|
||||
"elapsed_ms": elapsed_ms,
|
||||
**metrics,
|
||||
"error": "",
|
||||
}
|
||||
)
|
||||
except Exception as error: # Continue to expose the full failure matrix.
|
||||
row["error"] = str(error)
|
||||
rows.append(row)
|
||||
status = "ERROR" if row.get("error") else f"{row['actual_rate_pct']:.1f}% SSIM={row['ssim']:.4f}"
|
||||
print(f"[{case_number:02d}/{total:02d}] {input_path.name} -> {output_format} r{rate}: {status}", flush=True)
|
||||
if args.delay:
|
||||
time.sleep(args.delay)
|
||||
|
||||
summary = summarize(rows)
|
||||
report = {
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
"base_url": base_url,
|
||||
"inputs": len(inputs),
|
||||
"cases": len(rows),
|
||||
"errors": sum(bool(row.get("error")) for row in rows),
|
||||
"format_failures": sum(not bool(row.get("format_ok")) for row in rows if not row.get("error")),
|
||||
"target_failures": sum(not bool(row.get("target_met")) for row in rows if not row.get("error")),
|
||||
"summary": summary,
|
||||
"results": rows,
|
||||
}
|
||||
(output_dir / "report.json").write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
fieldnames = sorted({key for row in rows for key in row})
|
||||
with (output_dir / "results.csv").open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
print("\nformat rate target format-ok actual% saved% SSIM PSNR pixel% median-ms")
|
||||
for item in summary:
|
||||
print(
|
||||
f"{item['output_format']:>6} {item['requested_rate']:>4} "
|
||||
f"{item['target_met']}/{item['cases']} {item['format_ok']}/{item['cases']} "
|
||||
f"{item['mean_actual_rate_pct']:>7.2f} {item['mean_saved_pct']:>6.2f} "
|
||||
f"{item['mean_ssim']:.4f} {item['mean_psnr_db']:>5.2f} "
|
||||
f"{item['mean_pixel_ratio_pct']:>6.2f} {item['median_elapsed_ms']:>9.1f}"
|
||||
)
|
||||
print(f"\nreport: {output_dir / 'report.json'}")
|
||||
return 1 if args.strict and (report["errors"] or report["format_failures"] or report["target_failures"]) else 0
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
generate = subparsers.add_parser("generate", help="download and build the fixed real-photo corpus")
|
||||
generate.add_argument("--output-dir", default=".bench/corpus")
|
||||
generate.add_argument("--width", type=int, default=960)
|
||||
generate.add_argument("--height", type=int, default=640)
|
||||
|
||||
run = subparsers.add_parser("run", help="benchmark a running ImageForge deployment")
|
||||
run.add_argument("--base-url", default="http://127.0.0.1:8080")
|
||||
run.add_argument("--input-dir", default=".bench/corpus")
|
||||
run.add_argument("--output-dir", default=".bench/results")
|
||||
run.add_argument("--formats", default="jpeg,webp,avif")
|
||||
run.add_argument("--rates", default="30,50,70")
|
||||
run.add_argument("--email", default="")
|
||||
run.add_argument("--password", default="")
|
||||
run.add_argument("--token", default="")
|
||||
run.add_argument("--timeout", type=int, default=300)
|
||||
run.add_argument("--delay", type=float, default=0.05)
|
||||
run.add_argument("--strict", action="store_true")
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = build_parser().parse_args()
|
||||
if args.command == "generate":
|
||||
generate_corpus(Path(args.output_dir), args.width, args.height)
|
||||
return 0
|
||||
return benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
raise SystemExit(main())
|
||||
except KeyboardInterrupt:
|
||||
raise SystemExit(130) from None
|
||||
except Exception as error:
|
||||
print(f"benchmark failed: {error}", file=sys.stderr)
|
||||
raise SystemExit(1) from error
|
||||
Reference in New Issue
Block a user