Zhengshou Lai
dea83774a5
- --size now accepts free-form values (2K, 3K, 4K, 2304x1728, etc.) instead of a fixed [1K, 2K] choice. - Add --ratio shortcut (1:1, 4:3, 3:4, 16:9, 9:16) mapping to pixel sizes. - When both --size and --ratio are given, --size takes priority. - Default remains 2K when neither is specified.
252 lines
8.4 KiB
Python
252 lines
8.4 KiB
Python
"""Image utilities."""
|
||
|
||
import base64
|
||
import os
|
||
import subprocess
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
import click
|
||
|
||
from bin.config import config
|
||
from bin.utils import handle_errors, run_command
|
||
|
||
|
||
@click.group()
|
||
def image():
|
||
"""Image manipulation commands."""
|
||
pass
|
||
|
||
|
||
@image.command("eps-to-pdf")
|
||
@click.argument("files", nargs=-1, required=True)
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def eps_to_pdf(files, dry_run):
|
||
"""Convert EPS files to PDF."""
|
||
for pattern in files:
|
||
for eps_file in Path(".").glob(pattern):
|
||
if eps_file.suffix.lower() not in (".eps", ".ps"):
|
||
continue
|
||
output = eps_file.with_suffix(".pdf")
|
||
|
||
if dry_run:
|
||
click.echo(f"Would convert: {eps_file} -> {output}")
|
||
continue
|
||
|
||
click.echo(f"Converting: {eps_file} -> {output}")
|
||
run_command(["epstopdf", str(eps_file)], check=False)
|
||
|
||
|
||
@image.command("eps-fix")
|
||
@click.argument("file", type=click.Path(exists=True))
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def eps_fix(file, dry_run):
|
||
"""Fix EPS files created by print command."""
|
||
path = Path(file)
|
||
content = path.read_text()
|
||
|
||
if "/f/fill" not in content:
|
||
click.echo("File already fixed or not created by print command")
|
||
return
|
||
|
||
if dry_run:
|
||
click.echo(f"Would fix: {file}")
|
||
return
|
||
|
||
lines = content.split("\n")
|
||
new_lines = []
|
||
for line in lines:
|
||
if "/f/fill" in line:
|
||
line = line.replace("/f/fill", "")
|
||
new_lines.append(line)
|
||
|
||
path.write_text("\n".join(new_lines))
|
||
click.echo(f"Fixed: {file}")
|
||
|
||
|
||
@image.command("jpg-to-pdf")
|
||
@click.argument("files", nargs=-1, required=True)
|
||
@click.option("--output", "-o", help="Output PDF file")
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def jpg_to_pdf(files, output, dry_run):
|
||
"""Convert JPG files to PDF."""
|
||
images = []
|
||
for pattern in files:
|
||
for jpg_file in Path(".").glob(pattern):
|
||
if jpg_file.suffix.lower() in (".jpg", ".jpeg"):
|
||
images.append(str(jpg_file))
|
||
|
||
if not images:
|
||
click.echo("No JPG files found")
|
||
return
|
||
|
||
output = output or "output.pdf"
|
||
|
||
if dry_run:
|
||
click.echo(f"Would convert {len(images)} images to {output}")
|
||
return
|
||
|
||
click.echo(f"Converting {len(images)} images to {output}")
|
||
run_command(["convert"] + images + [output], check=False)
|
||
|
||
|
||
@image.command("tiff-to-pdf")
|
||
@click.argument("files", nargs=-1, required=True)
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def tiff_to_pdf(files, dry_run):
|
||
"""Convert TIFF files to PDF."""
|
||
for pattern in files:
|
||
for tiff_file in Path(".").glob(pattern):
|
||
if tiff_file.suffix.lower() not in (".tif", ".tiff"):
|
||
continue
|
||
output = tiff_file.with_suffix(".pdf")
|
||
|
||
if dry_run:
|
||
click.echo(f"Would convert: {tiff_file} -> {output}")
|
||
continue
|
||
|
||
click.echo(f"Converting: {tiff_file} -> {output}")
|
||
run_command(["convert", str(tiff_file), str(output)], check=False)
|
||
|
||
|
||
@image.command("tiff-compress")
|
||
@click.argument("files", nargs=-1, required=True)
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def tiff_compress(files, dry_run):
|
||
"""Compress TIFF files."""
|
||
for pattern in files:
|
||
for tiff_file in Path(".").glob(pattern):
|
||
if tiff_file.suffix.lower() not in (".tif", ".tiff"):
|
||
continue
|
||
|
||
if dry_run:
|
||
click.echo(f"Would compress: {tiff_file}")
|
||
continue
|
||
|
||
click.echo(f"Compressing: {tiff_file}")
|
||
run_command(
|
||
["convert", str(tiff_file), "-compress", "zip", str(tiff_file)],
|
||
check=False,
|
||
)
|
||
|
||
|
||
@image.command("compress")
|
||
@click.argument("files", nargs=-1, required=True)
|
||
@click.option("--target-mb", "-t", default=1.0, help="Target max file size in MB")
|
||
@click.option("--max-width", "-w", default=1600, help="Max width in pixels if resize needed")
|
||
@click.option("--quality", "-q", default=75, help="JPEG quality for resized images")
|
||
@click.option("--suffix", "-s", default="_compressed", help="Output filename suffix")
|
||
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
||
@handle_errors
|
||
def compress_images(files, target_mb, max_width, quality, suffix, dry_run):
|
||
"""Compress PNG/JPG images to fit under target size."""
|
||
from PIL import Image
|
||
|
||
for pattern in files:
|
||
for img_file in Path(".").glob(pattern):
|
||
if img_file.suffix.lower() not in (".png", ".jpg", ".jpeg"):
|
||
continue
|
||
|
||
output = img_file.with_stem(f"{img_file.stem}{suffix}")
|
||
if img_file.suffix.lower() == ".png":
|
||
output = output.with_suffix(".jpg")
|
||
|
||
if dry_run:
|
||
click.echo(f"Would compress: {img_file} -> {output}")
|
||
continue
|
||
|
||
img = Image.open(img_file)
|
||
if img.mode in ("RGBA", "P"):
|
||
img = img.convert("RGB")
|
||
|
||
# Try full-res first with q80
|
||
img.save(output, "JPEG", quality=80, optimize=True)
|
||
size_mb = output.stat().st_size / (1024 * 1024)
|
||
|
||
if size_mb > target_mb:
|
||
w, h = img.size
|
||
if w > max_width:
|
||
h = int(h * max_width / w)
|
||
w = max_width
|
||
try:
|
||
resample = Image.Resampling.LANCZOS
|
||
except AttributeError:
|
||
resample = Image.LANCZOS # type: ignore[attr-defined]
|
||
img = img.resize((w, h), resample)
|
||
img.save(output, "JPEG", quality=quality, optimize=True)
|
||
size_mb = output.stat().st_size / (1024 * 1024)
|
||
if size_mb > target_mb:
|
||
img.save(output, "JPEG", quality=quality - 5, optimize=True)
|
||
size_mb = output.stat().st_size / (1024 * 1024)
|
||
|
||
click.echo(f"Compressed: {img_file} -> {output} ({size_mb:.2f} MB)")
|
||
|
||
|
||
_RATIO_MAP = {
|
||
"1:1": "1024x1024",
|
||
"4:3": "1024x768",
|
||
"3:4": "768x1024",
|
||
"16:9": "1280x720",
|
||
"9:16": "720x1280",
|
||
}
|
||
|
||
|
||
@image.command("generate")
|
||
@click.argument("prompt", required=False)
|
||
@click.option("-f", "--file", type=click.Path(exists=True), help="从文件读取 prompt")
|
||
@click.option("-s", "--size", help='图片尺寸 (如 2K, 3K, 4K 或 2304x1728)')
|
||
@click.option("-r", "--ratio", type=click.Choice(list(_RATIO_MAP.keys())), help="快捷比例 (1:1, 4:3, 3:4, 16:9, 9:16)")
|
||
@click.option("-o", "--output", help="保存路径（默认只输出 URL）")
|
||
@click.option("--no-watermark", is_flag=True, help="不添加豆包水印")
|
||
@click.option("--b64", is_flag=True, help="使用 b64_json 格式获取图片数据")
|
||
@handle_errors
|
||
def generate_image_cmd(prompt, file, size, ratio, output, no_watermark, b64):
|
||
"""使用豆包/Seedream 模型生成图片。"""
|
||
from openai import OpenAI
|
||
|
||
if file:
|
||
prompt = Path(file).read_text(encoding="utf-8").strip()
|
||
elif not prompt:
|
||
raise click.UsageError("必须提供 prompt 或使用 -f/--file 从文件读取")
|
||
|
||
api_key = config.get("apikey_ark")
|
||
if not api_key:
|
||
click.echo("Error: apikey_ark not set. Run: mytoolkit env set apikey_ark <value>", err=True)
|
||
raise click.Abort()
|
||
|
||
if size and ratio:
|
||
click.echo("Warning: --size 和 --ratio 同时指定，--size 优先", err=True)
|
||
|
||
resolved_size = size or (ratio and _RATIO_MAP[ratio]) or "2K"
|
||
|
||
client = OpenAI(
|
||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||
api_key=api_key,
|
||
)
|
||
|
||
resp = client.images.generate(
|
||
model="doubao-seedream-5-0-260128",
|
||
prompt=prompt,
|
||
size=resolved_size,
|
||
response_format="b64_json" if b64 else "url",
|
||
extra_body={"watermark": not no_watermark},
|
||
)
|
||
|
||
if b64:
|
||
image_bytes = base64.b64decode(resp.data[0].b64_json)
|
||
save_path = output or "generated_image.png"
|
||
Path(save_path).write_bytes(image_bytes)
|
||
click.echo(os.path.abspath(save_path))
|
||
else:
|
||
url = resp.data[0].url
|
||
if output:
|
||
urllib.request.urlretrieve(url, output)
|
||
click.echo(os.path.abspath(output))
|
||
else:
|
||
click.echo(url)
|