"""
Anime Style Transfer — FastAPI service
========================================
Pure FastAPI (no Gradio). Exposes:
POST /anime -> upload an image, get back a stylized PNG
GET /health -> simple liveness check
GET /docs -> auto-generated Swagger UI to test the API in-browser
GET / -> minimal HTML test page (manual UI, no extra deps)
Model (AnimeGANv2) loads ONCE at startup, not per-request, so inference
stays fast across repeated calls.
"""
import io
import cv2
import numpy as np
import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import HTMLResponse, StreamingResponse
from PIL import Image, ImageEnhance, ImageFilter
from torchvision.transforms.functional import to_pil_image, to_tensor
# ---------------------------------------------------------------------------
# Model setup
# ---------------------------------------------------------------------------
DEVICE = "cpu"
torch.set_num_threads(4) # tune to your container's CPU allocation
STYLES = {
"face_v2": "face_paint_512_v2",
"face_v1": "face_paint_512_v1",
"celeb": "celeba_distill",
"paprika": "paprika",
}
DEFAULT_STYLE = "face_v2"
_model_cache: dict = {}
# Path where the Dockerfile clones the model repo at build time. Loading from
# here with source="local" skips torch.hub's GitHub API call entirely (the
# call that was hitting GitHub's rate limit and crashing startup with a 403).
_LOCAL_REPO_PATH = "/app/animegan2-pytorch"
def load_model(style_key: str):
if style_key not in STYLES:
raise HTTPException(
status_code=400,
detail=f"Unknown style '{style_key}'. Valid options: {list(STYLES.keys())}",
)
if style_key not in _model_cache:
model = torch.hub.load(
_LOCAL_REPO_PATH,
"generator",
source="local",
pretrained=STYLES[style_key],
device=DEVICE,
)
model.eval()
_model_cache[style_key] = model
return _model_cache[style_key]
def _resize_preserving_aspect(image: Image.Image, target_long_side: int, multiple: int = 32) -> Image.Image:
"""
Resize so the longer side equals target_long_side, keeping aspect ratio,
then round both dimensions down to a multiple of `multiple` (AnimeGAN's
conv layers expect dimensions divisible by 32 for clean upsampling).
Forcing a square here is what caused warped/blurry output on non-square
photos before.
"""
w, h = image.size
scale = target_long_side / max(w, h)
new_w, new_h = round(w * scale), round(h * scale)
new_w = max(multiple, (new_w // multiple) * multiple)
new_h = max(multiple, (new_h // multiple) * multiple)
return image.resize((new_w, new_h), Image.LANCZOS)
# Loaded once — ships with opencv-python, no extra download needed.
_FACE_CASCADE = cv2.CascadeClassifier(
cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
)
def detect_and_crop_face(image: Image.Image, margin: float = 0.6) -> Image.Image | None:
"""
Find the largest face in the image and crop around it with a margin so
eyes/mouth occupy a bigger share of the pixels the GAN sees. This is the
main lever for sharper eyes, since AnimeGANv2 struggles most on small
high-detail regions — making them bigger in the input helps the most.
Returns None if no face is found, so the caller can fall back to using
the full image untouched (important for non-portrait photos).
"""
cv_img = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
faces = _FACE_CASCADE.detectMultiScale(
gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60)
)
if len(faces) == 0:
return None
# Largest detected face — most likely the main subject
x, y, w, h = max(faces, key=lambda f: f[2] * f[3])
# Expand the box by `margin` on each side so we keep forehead/chin/hair,
# not just a tight box that would clip the face awkwardly.
img_w, img_h = image.size
pad_x, pad_y = int(w * margin), int(h * margin)
left = max(0, x - pad_x)
top = max(0, y - pad_y)
right = min(img_w, x + w + pad_x)
bottom = min(img_h, y + h + pad_y)
return image.crop((left, top, right, bottom))
def sharpen_output(image: Image.Image, amount: float = 1.4) -> Image.Image:
"""
Unsharp mask pass on the stylized result. This crisps up edges the GAN
already drew (eye outlines, lash lines, line art) — it cannot invent
detail that wasn't generated, but it noticeably reduces the "soft/smudged"
look AnimeGANv2 tends to produce, especially around eyes.
"""
return image.filter(
ImageFilter.UnsharpMask(radius=2, percent=int(100 * amount), threshold=2)
)
def generate_stylized_image(image: Image.Image, style_key: str, size: int, face_crop: bool = False) -> Image.Image:
"""
STAGE 1 — generation.
Runs the GAN and returns the raw stylized output, unenhanced.
"""
model = load_model(style_key)
source = image.convert("RGB")
if face_crop:
cropped = detect_and_crop_face(source)
if cropped is not None:
source = cropped
# if no face found, silently fall back to the full image
img = _resize_preserving_aspect(source, target_long_side=size)
inp = to_tensor(img).unsqueeze(0) * 2 - 1
with torch.no_grad():
out = model(inp) # blocking call — by the time this returns, output is ready
raw_output = to_pil_image((out.squeeze() * 0.5 + 0.5).clamp(0, 1))
return raw_output
def enhance_image(image: Image.Image, sharpen_amount: float = 1.4, boost_color: bool = True) -> Image.Image:
"""
STAGE 2 — enhancement.
Takes a completed (ready) stylized image and improves perceived clarity:
1. Unsharp mask — crisps up edges already drawn by the GAN (eye
outlines, lash lines, line art) without inventing new detail.
2. Mild contrast + color boost — AnimeGAN output tends to look a touch
flat/washed; a small boost gives it a more typical "anime" punch.
Only called once Stage 1's output is fully ready, never interleaved
with generation.
"""
sharpened = sharpen_output(image, amount=sharpen_amount)
if boost_color:
sharpened = ImageEnhance.Contrast(sharpened).enhance(1.08)
sharpened = ImageEnhance.Color(sharpened).enhance(1.12)
return sharpened
def run_inference(
image: Image.Image,
style_key: str,
size: int,
face_crop: bool = False,
enhance: bool = True,
) -> Image.Image:
"""
Full pipeline, run strictly in order:
1. Generate the stylized image (Stage 1).
2. Once that output is ready, optionally enhance it (Stage 2).
3. Return the final output.
"""
# --- Stage 1: generation ---
raw_output = generate_stylized_image(image, style_key, size, face_crop=face_crop)
# --- Stage 2: enhancement (only runs on the completed Stage 1 output) ---
final_output = enhance_image(raw_output) if enhance else raw_output
return final_output
# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------
app = FastAPI(
title="Anime Style Transfer API",
description="Upload a photo, get back an anime-styled version. Powered by AnimeGANv2.",
version="1.0.0",
)
@app.on_event("startup")
def warm_up():
# Load the default model at startup so the first real request is fast,
# not slowed down by a cold model load.
load_model(DEFAULT_STYLE)
@app.get("/health")
def health():
return {"status": "ok"}
@app.post(
"/anime",
summary="Convert a photo to anime style",
response_description="PNG image of the stylized result",
)
async def anime_endpoint(
file: UploadFile = File(..., description="Image file (jpg/png) to convert"),
style: str = DEFAULT_STYLE,
size: int = 512,
face_crop: bool = False,
enhance: bool = True,
):
"""
- **file**: the photo to convert (multipart upload)
- **style**: one of `face_v2`, `face_v1`, `celeb`, `paprika`
- **size**: square resolution for inference (lower = faster). 512 is a
good default; try 384 or 256 if you need more speed on CPU.
- **face_crop**: auto-detect and crop to the main face before stylizing.
Makes eyes/mouth occupy more pixels, which usually improves facial
detail. Falls back to the full image if no face is detected.
- **enhance**: apply a sharpening pass to the output to reduce the
soft/smudged look around fine details like eyes.
"""
if not file.content_type or not file.content_type.startswith("image/"):
raise HTTPException(status_code=400, detail="Uploaded file must be an image.")
if size < 128 or size > 1024:
raise HTTPException(status_code=400, detail="size must be between 128 and 1024.")
try:
raw = await file.read()
image = Image.open(io.BytesIO(raw))
except Exception:
raise HTTPException(status_code=400, detail="Could not read the uploaded file as an image.")
result = run_inference(image, style, size, face_crop=face_crop, enhance=enhance)
buf = io.BytesIO()
result.save(buf, format="PNG")
buf.seek(0)
# Confirm what actually happened to produce this output — lets a caller
# verify the final image was enhanced rather than assuming it from the
# request params alone (e.g. if face_crop silently fell back because no
# face was found, or enhance was forced off by a future code change).
headers = {
"X-Pipeline-Style": style,
"X-Pipeline-Enhanced": "true" if enhance else "false",
"X-Pipeline-Face-Crop-Requested": "true" if face_crop else "false",
}
return StreamingResponse(buf, media_type="image/png", headers=headers)
# ---------------------------------------------------------------------------
# Minimal built-in test UI — no Gradio, just plain HTML/JS served by FastAPI
# ---------------------------------------------------------------------------
_TEST_PAGE = """
Anime Style Transfer
Anime StyleTransfer
Upload a photo, pick a style, get anime art back. Full API schema at /docs.