""" Image Compressor Pro — Python Backend (FastAPI + Pillow) ======================================================== A high-performance, concurrent image compression API designed to replace the Node.js + Sharp backend. Uses Pillow (with pillow-heif for HEIF/AVIF) and FastAPI with StreamingResponse for efficient, non-blocking I/O. Deployment target: HuggingFace Spaces (Docker) or local testing. """ import io import logging from contextlib import asynccontextmanager from concurrent.futures import ThreadPoolExecutor from fastapi import FastAPI, Request, Query, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse, PlainTextResponse from starlette.concurrency import run_in_threadpool from PIL import Image # --------------------------------------------------------------------------- # Optional: HEIF / AVIF support via pillow-heif # --------------------------------------------------------------------------- try: import pillow_heif pillow_heif.register_heif_opener() # enables Image.open() for HEIF/AVIF HEIF_AVAILABLE = True except ImportError: HEIF_AVAILABLE = False # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)-7s | %(message)s", datefmt="%H:%M:%S", ) log = logging.getLogger("compressor") # --------------------------------------------------------------------------- # Thread pool — used to offload CPU-bound Pillow work # --------------------------------------------------------------------------- MAX_WORKERS = 4 _pool = ThreadPoolExecutor(max_workers=MAX_WORKERS) # --------------------------------------------------------------------------- # Lifespan (startup / shutdown) # --------------------------------------------------------------------------- @asynccontextmanager async def lifespan(app: FastAPI): log.info("Compressor backend starting (workers=%d)", MAX_WORKERS) if HEIF_AVAILABLE: log.info("pillow-heif is available — HEIF/AVIF support enabled") else: log.warning("pillow-heif not installed — HEIF/AVIF output disabled") yield _pool.shutdown(wait=False) log.info("Compressor backend shut down") # --------------------------------------------------------------------------- # App # --------------------------------------------------------------------------- app = FastAPI( title="Image Compressor Pro API", version="1.0.0", lifespan=lifespan, ) # CORS — allow all origins so the frontend can call from any domain app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- MAX_FILE_SIZE = 20 * 1024 * 1024 # 20 MB # MIME types for each supported output format MIME_MAP = { "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp", "avif": "image/avif", "heif": "image/heic", "tiff": "image/tiff", "gif": "image/gif", "bmp": "image/bmp", "ico": "image/x-icon", "jp2": "image/jp2", } # --------------------------------------------------------------------------- # Core compression logic (runs in thread pool) # --------------------------------------------------------------------------- def _compress_image(raw_bytes: bytes, quality: int, fmt: str) -> bytes: """ Open *raw_bytes* as an image, compress it into *fmt* at the given *quality*, and return the resulting bytes. This is intentionally a **synchronous** function because Pillow is CPU-bound. It is called via ``run_in_threadpool`` so the event loop is never blocked. """ img = Image.open(io.BytesIO(raw_bytes)) # Convert RGBA → RGB for formats that don't support alpha if fmt in ("jpeg", "tiff", "bmp", "jp2") and img.mode in ("RGBA", "LA", "PA"): background = Image.new("RGB", img.size, (255, 255, 255)) background.paste(img, mask=img.split()[-1]) # alpha composite img = background elif fmt in ("jpeg", "bmp", "jp2") and img.mode not in ("RGB", "L"): img = img.convert("RGB") elif fmt == "gif" and img.mode == "RGBA": # GIF supports palette transparency; convert RGBA → P with transparency img = img.convert("RGBA") # ensure consistent mode alpha = img.split()[-1] img = img.convert("RGB").convert("P", palette=Image.ADAPTIVE, colors=255) mask = Image.eval(alpha, lambda a: 255 if a <= 128 else 0) img.paste(255, mask) # map transparent pixels to index 255 # transparency index will be set during save out_buffer = io.BytesIO() if fmt == "jpeg": img.save( out_buffer, format="JPEG", quality=quality, optimize=True, progressive=True, subsampling="4:2:0" if quality < 90 else "4:4:4", ) elif fmt == "png": # PNG quality mapped to compress_level (0-9, higher = smaller) compress_level = max(1, min(9, 9 - (quality // 12))) img.save( out_buffer, format="PNG", compress_level=compress_level, optimize=True, ) elif fmt == "webp": img.save( out_buffer, format="WEBP", quality=quality, method=4, # 0-6, higher = slower + better compression ) elif fmt == "avif": # NOTE: Pillow 12.x has NATIVE AVIF support via AvifImagePlugin. # It does NOT use pillow-heif's 'chroma' parameter. # The correct parameter is 'subsampling' (default '4:2:0'). # We MUST use '4:4:4' to prevent alpha edge bleeding on transparent images. if img.mode in ("RGBA", "LA", "PA") or "transparency" in img.info: img = img.convert("RGBA") import numpy as np data = np.array(img) alpha = data[:, :, 3] # Remove hidden RGB in fully transparent pixels data[:, :, :3][alpha == 0] = [0, 0, 0] img = Image.fromarray(data, "RGBA") elif img.mode not in ("RGB", "L"): img = img.convert("RGB") img.save( out_buffer, format="AVIF", quality=quality, subsampling="4:4:4", ) elif fmt == "heif": if not HEIF_AVAILABLE: raise ValueError("HEIF output requires pillow-heif (not installed)") img.save( out_buffer, format="HEIF", quality=quality, ) elif fmt == "tiff": img.save( out_buffer, format="TIFF", compression="tiff_lzw", # lossless LZW compression ) elif fmt == "gif": # Preserve animated GIF frames if present n_frames = getattr(img, 'n_frames', 1) if n_frames > 1: frames = [] durations = [] for i in range(n_frames): img.seek(i) frame = img.copy() if frame.mode != "P": frame = frame.convert("RGBA").convert("P", palette=Image.ADAPTIVE, colors=256) frames.append(frame) durations.append(img.info.get('duration', 100)) frames[0].save( out_buffer, format="GIF", save_all=True, append_images=frames[1:], duration=durations, loop=img.info.get('loop', 0), optimize=True, ) else: if img.mode not in ("P", "L"): img = img.convert("P", palette=Image.ADAPTIVE, colors=256) img.save( out_buffer, format="GIF", optimize=True, ) elif fmt == "bmp": if img.mode not in ("RGB", "L", "1"): img = img.convert("RGB") img.save( out_buffer, format="BMP", ) elif fmt == "ico": # ICO spec limits dimensions to 256×256 max_ico = 256 if img.width > max_ico or img.height > max_ico: img.thumbnail((max_ico, max_ico), Image.LANCZOS) if img.mode != "RGBA": img = img.convert("RGBA") img.save( out_buffer, format="ICO", ) elif fmt == "jp2": if img.mode not in ("RGB", "L", "RGBA"): img = img.convert("RGB") img.save( out_buffer, format="JPEG2000", quality_mode="rates", quality_layers=[quality], ) else: raise ValueError(f"Unsupported format: {fmt}") out_buffer.seek(0) return out_buffer.getvalue() # --------------------------------------------------------------------------- # Routes # --------------------------------------------------------------------------- @app.get("/ping", response_class=PlainTextResponse) async def health_check(): """Simple health-check endpoint compatible with the Node backend.""" return "pong" @app.get("/", response_class=PlainTextResponse) async def root(): """Root endpoint — useful for HuggingFace Spaces health probes.""" return "Image Compressor Pro API is running" @app.post("/api/process-stream") async def process_stream( request: Request, quality: int = Query(default=80, ge=1, le=100), format: str = Query(default="jpeg", pattern="^(jpeg|png|webp|avif|heif|tiff|gif|bmp|ico|jp2)$"), ): """ Receive a raw image in the request body (Content-Type: image/*), compress it using Pillow, and stream the result back. Query parameters ---------------- quality : int (1–100) Compression quality. Higher = better quality, larger file. format : str Target output format (jpeg, png, webp, avif, heif, tiff). This mirrors the Node.js backend's ``POST /api/process-stream`` contract so the frontend works without changes. """ # ---- Read the incoming image bytes (with size guard) ---- body = await request.body() if len(body) == 0: raise HTTPException(status_code=400, detail="Empty request body") if len(body) > MAX_FILE_SIZE: raise HTTPException( status_code=413, detail=f"File too large ({len(body)} bytes). Max is {MAX_FILE_SIZE} bytes.", ) # ---- Check HEIF/AVIF availability ---- if format in ("avif", "heif") and not HEIF_AVAILABLE: raise HTTPException( status_code=400, detail=f"{format.upper()} format requires pillow-heif which is not installed.", ) # ---- Offload CPU-bound compression to thread pool ---- try: compressed = await run_in_threadpool(_compress_image, body, quality, format) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) except Exception as e: log.exception("Compression failed") raise HTTPException(status_code=500, detail=f"Compression error: {e}") # ---- Stream result back ---- mime = MIME_MAP.get(format, "application/octet-stream") return StreamingResponse( io.BytesIO(compressed), media_type=mime, headers={ "Content-Disposition": f"attachment; filename=compressed_image.{format}", "Content-Length": str(len(compressed)), "X-Original-Size": str(len(body)), "X-Compressed-Size": str(len(compressed)), }, ) # --------------------------------------------------------------------------- # Entry point for local development # --------------------------------------------------------------------------- if __name__ == "__main__": import uvicorn uvicorn.run( "app:app", host="0.0.0.0", port=7860, workers=1, # single worker for local dev; increase for prod log_level="info", reload=True, # auto-reload on code changes during dev )