71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
from fastapi import HTTPException, status
|
|
from PIL import Image, ImageFilter
|
|
|
|
from ..config import Settings
|
|
from ..image_io import DownloadedImage, load_normalized_image
|
|
from .base import EngineHealth, UpscaleEngine, UpscaleResult
|
|
|
|
|
|
MODE_PROFILES = {
|
|
"standard": {"profile": "general", "sharpen_percent": 120, "radius": 1.0, "threshold": 3},
|
|
"artwork": {"profile": "artwork", "sharpen_percent": 150, "radius": 1.2, "threshold": 2},
|
|
"photo": {"profile": "photo", "sharpen_percent": 95, "radius": 0.8, "threshold": 4},
|
|
"illustration": {"profile": "illustration", "sharpen_percent": 135, "radius": 1.0, "threshold": 2},
|
|
}
|
|
|
|
|
|
class PillowUpscaleEngine(UpscaleEngine):
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
|
|
def health(self) -> EngineHealth:
|
|
return EngineHealth(
|
|
status="ok",
|
|
engine="pillow",
|
|
device=self.settings.device,
|
|
models_loaded=True,
|
|
)
|
|
|
|
def upscale(self, downloaded: DownloadedImage, scale: int, mode: str, output_format: str) -> UpscaleResult:
|
|
started_at = time.perf_counter()
|
|
profile = MODE_PROFILES[mode]
|
|
image = load_normalized_image(downloaded.path)
|
|
width, height = image.size
|
|
target_width = width * scale
|
|
target_height = height * scale
|
|
|
|
if target_width > self.settings.max_output_width or target_height > self.settings.max_output_height:
|
|
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Worker rejected the image.")
|
|
|
|
result = image.resize((target_width, target_height), Image.Resampling.LANCZOS)
|
|
result = result.filter(
|
|
ImageFilter.UnsharpMask(
|
|
radius=profile["radius"],
|
|
percent=profile["sharpen_percent"],
|
|
threshold=profile["threshold"],
|
|
)
|
|
)
|
|
|
|
return UpscaleResult(
|
|
image=result,
|
|
metadata={
|
|
"engine": "pillow",
|
|
"model": "pillow-lanczos",
|
|
"requested_scale": scale,
|
|
"native_model_scale": scale,
|
|
"mode": mode,
|
|
"device": self.settings.device,
|
|
"profile": profile["profile"],
|
|
"real_ai_upscale": False,
|
|
"processing_seconds": round(time.perf_counter() - started_at, 3),
|
|
"input_width": width,
|
|
"input_height": height,
|
|
"output_width": target_width,
|
|
"output_height": target_height,
|
|
"output_format": output_format,
|
|
},
|
|
) |