98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
|
"""File I/O utilities: base64 encoding, output path generation, download."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import base64
|
||
|
|
from datetime import datetime
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from rich.console import Console
|
||
|
|
from rich.progress import (
|
||
|
|
BarColumn,
|
||
|
|
DownloadColumn,
|
||
|
|
Progress,
|
||
|
|
SpinnerColumn,
|
||
|
|
TimeElapsedColumn,
|
||
|
|
TransferSpeedColumn,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def image_to_base64(path: Path) -> str:
|
||
|
|
"""Read an image file and return a base64-encoded string."""
|
||
|
|
suffix = path.suffix.lower().lstrip(".")
|
||
|
|
mime_map = {
|
||
|
|
"jpg": "image/jpeg",
|
||
|
|
"jpeg": "image/jpeg",
|
||
|
|
"png": "image/png",
|
||
|
|
"gif": "image/gif",
|
||
|
|
"webp": "image/webp",
|
||
|
|
}
|
||
|
|
mime = mime_map.get(suffix, "image/jpeg")
|
||
|
|
with open(path, "rb") as f:
|
||
|
|
encoded = base64.b64encode(f.read()).decode()
|
||
|
|
return f"data:{mime};base64,{encoded}"
|
||
|
|
|
||
|
|
|
||
|
|
def video_to_base64(path: Path) -> str:
|
||
|
|
"""Read a video file and return a base64-encoded string."""
|
||
|
|
suffix = path.suffix.lower().lstrip(".")
|
||
|
|
mime_map = {
|
||
|
|
"mp4": "video/mp4",
|
||
|
|
"mov": "video/quicktime",
|
||
|
|
"avi": "video/x-msvideo",
|
||
|
|
"webm": "video/webm",
|
||
|
|
}
|
||
|
|
mime = mime_map.get(suffix, "video/mp4")
|
||
|
|
with open(path, "rb") as f:
|
||
|
|
encoded = base64.b64encode(f.read()).decode()
|
||
|
|
return f"data:{mime};base64,{encoded}"
|
||
|
|
|
||
|
|
|
||
|
|
def auto_output_path(task_type: str, model: str, ext: str = "jpg", output_dir: str = ".") -> Path:
|
||
|
|
"""Generate a timestamped output filename."""
|
||
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
|
|
safe_model = model.replace("/", "-").replace(" ", "-")
|
||
|
|
filename = f"freepik_{task_type}_{safe_model}_{ts}.{ext}"
|
||
|
|
return Path(output_dir) / filename
|
||
|
|
|
||
|
|
|
||
|
|
def save_from_url(url: str, dest: Path, console: Console) -> None:
|
||
|
|
"""Stream-download a URL to a file, displaying a Rich progress bar."""
|
||
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
with Progress(
|
||
|
|
SpinnerColumn(style="bold magenta"),
|
||
|
|
"[progress.description]{task.description}",
|
||
|
|
BarColumn(bar_width=30, style="magenta", complete_style="green"),
|
||
|
|
"[progress.percentage]{task.percentage:>3.0f}%",
|
||
|
|
DownloadColumn(),
|
||
|
|
TransferSpeedColumn(),
|
||
|
|
TimeElapsedColumn(),
|
||
|
|
console=console,
|
||
|
|
transient=True,
|
||
|
|
) as progress:
|
||
|
|
with httpx.stream("GET", url, follow_redirects=True, timeout=120) as r:
|
||
|
|
r.raise_for_status()
|
||
|
|
total = int(r.headers.get("content-length", 0)) or None
|
||
|
|
task = progress.add_task(
|
||
|
|
f"[dim]Saving[/dim] [bold]{dest.name}[/bold]",
|
||
|
|
total=total,
|
||
|
|
)
|
||
|
|
with open(dest, "wb") as f:
|
||
|
|
for chunk in r.iter_bytes(chunk_size=65536):
|
||
|
|
f.write(chunk)
|
||
|
|
progress.advance(task, len(chunk))
|
||
|
|
|
||
|
|
|
||
|
|
def get_image_dimensions(path: Path) -> tuple[Optional[int], Optional[int]]:
|
||
|
|
"""Return (width, height) of an image using Pillow."""
|
||
|
|
try:
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
with Image.open(path) as img:
|
||
|
|
return img.width, img.height
|
||
|
|
except Exception:
|
||
|
|
return None, None
|