""" upscale.py — POST /upscale/upload and GET /upscale/result/{job_id} INDEPENDENT of the LaMa/watermark-removal pipeline. Do not import from jobs_store.py, schemas.py, generate.py, inpaint_engine.py, or any other watermark-removal module. Design ------ Unlike watermark removal (which separates upload from inference because it needs a user-drawn mask between the two steps), upscaling has no user input after the image is chosen — so upload and inference kick-off are combined into a single POST /upscale/upload endpoint. Flow ---- 1. POST /upscale/upload (multipart file) a. Validate content-type. b. EXIF-correct and convert to RGB via Pillow (same reason as /upload). c. Assign UUID job_id; create storage/upscale/{job_id}/. d. Save EXIF-corrected image as original.jpg. e. Persist UpscaleJob(status="pending") to upscale_jobs.json. f. Immediately advance status to "processing" and enqueue BackgroundTask. g. Return {"job_id": job_id} immediately (inference runs in background). 2. GET /upscale/result/{job_id} (polling) Returns {"job_id", "status", "result_image_url", "error_message"}. result_image_url is a /storage/... URL served by the existing StaticFiles mount in main.py — no extra route needed. 3. run_upscale_job (background worker) Runs after the POST response is sent; has no access to the Request object. Receives the engine as an explicit argument (same pattern as run_inpaint_job). Updates job status to "completed" or "failed"; never raises. Storage layout (kept entirely separate from watermark-removal storage) ---------------------------------------------------------------------- storage/ upscale/ {job_id}/ original.jpg — EXIF-corrected upload result.jpg — 4× upscaled output (written on completion) These are served as: /storage/upscale/{job_id}/original.jpg /storage/upscale/{job_id}/result.jpg via the existing app.mount("/storage", StaticFiles(...)) in main.py. """ from __future__ import annotations import io import logging import math import time import traceback import uuid from concurrent.futures import ThreadPoolExecutor, TimeoutError as _FuturesTimeout from datetime import datetime, timezone from pathlib import Path import cv2 import numpy as np from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, Request, UploadFile from PIL import Image, ImageOps from app.config import STORAGE_DIR, UPSCALE_TILE_THRESHOLD, UPSCALE_TILE_SIZE from app.utils import job_folder_name from app.upscale_jobs_store import UpscaleJob, create_job, get_job, update_job from app.services.upscale_engine import BaseUpscaleEngine logger = logging.getLogger(__name__) router = APIRouter() # --------------------------------------------------------------------------- # Job-level timeout # --------------------------------------------------------------------------- # Hard ceiling on how long engine.upscale() is allowed to run. # # Justification (33 s/tile measured baseline on this machine): # A 2048×1367 image → 4×3 = 12 tiles → expected ~6.6 min at baseline. # The reported runaway job hit 36+ min (≥5× baseline) before manual kill. # 30 minutes gives 4.5× headroom over the 12-tile baseline — generous # enough for moderate thermal throttling — while cutting off a stuck job # before it can run indefinitely. At normal speed a practical worst-case # upload (~20 tiles for a very large image) takes ~11 min baseline; # 30 min = 2.7× headroom, still safe. # # Note: on timeout the *upscale thread* becomes a zombie — Python cannot # forcibly terminate it. It will continue consuming CPU in the background # until the computation finishes or the server restarts. The job record is # immediately marked "failed" so the user is not left waiting indefinitely. _UPSCALE_TIMEOUT_SECS: int = 30 * 60 # 30 minutes # Upscale jobs live in a dedicated sub-directory of storage/ so they can # never collide with watermark-removal job folders (which use storage/{uuid}/steps/). _UPSCALE_STORAGE = STORAGE_DIR / "upscale" # --------------------------------------------------------------------------- # Post-processing helpers # --------------------------------------------------------------------------- def _denoise(img: np.ndarray) -> np.ndarray: """Non-local means denoising (BGR uint8 in → BGR uint8 out). Averages similar patches across the image to reduce random noise and JPEG compression artifacts while preserving edges better than a simple blur. Parameters (h=10, hColor=10, templateWindowSize=7, searchWindowSize=21) are the established standard defaults for general-purpose color denoising. Applied BEFORE upscaling so we process the smaller original image — cheaper and removes artifacts before the model can amplify them 4×. """ return cv2.fastNlMeansDenoisingColored( img, None, h=10, hColor=10, templateWindowSize=7, searchWindowSize=21, ) def _sharpen(img: np.ndarray) -> np.ndarray: """Mild unsharp mask (RGB uint8 in → RGB uint8 out). Intentionally conservative (amount=0.4, sigma=1.0): the 4× Real-ESRGAN output already has an AI-oversharpened look on organic textures (observed during project testing — grass/foliage artifacts), so an aggressive sharpen would compound that problem. GaussianBlur and the arithmetic are channel-independent so no BGR conversion is needed. Applied AFTER upscaling, BEFORE any output_scale downscale step. """ blurred = cv2.GaussianBlur(img, (0, 0), sigmaX=1.0) img_f = img.astype(np.float32) blur_f = blurred.astype(np.float32) return np.clip(img_f + 0.4 * (img_f - blur_f), 0, 255).astype(np.uint8) # --------------------------------------------------------------------------- # Background worker # --------------------------------------------------------------------------- def run_upscale_job( job_id: str, folder_name: str, original_path: Path, engine: BaseUpscaleEngine, output_scale: str = "4x", denoise: bool = False, sharpen: bool = False, ) -> None: """Run Real-ESRGAN upscaling and save result; update job status. Called by FastAPI BackgroundTasks after the /upscale/upload response is sent. Has no access to the Request object — engine is passed as an explicit arg. Updates status to "completed" or "failed"; never raises. Pipeline order -------------- original → [denoise] → engine.upscale() → [sharpen] → [output_scale resize] → save denoise is applied on the small original image (cheaper; removes JPEG artifacts before the model can amplify them 4×). sharpen is applied on the full 4× output before any downscale, so the mild enhancement survives LANCZOS resampling. output_scale ------------ "4x" — save the full 4x output from the engine (current default behaviour). "2x" — downsample the 4x output by exactly 50% in each dimension using LANCZOS before saving. "1.5x"— downsample to 1.5× the original (resize factor 3/8 from 4× output). """ job_dir = _UPSCALE_STORAGE / folder_name result_path = job_dir / "result.jpg" t_job_start = time.monotonic() try: pil_img = Image.open(original_path).convert("RGB") img_w, img_h = pil_img.size # Log job dimensions and expected tile count so we can rule out a # tile-count bug when diagnosing slow jobs. if max(img_w, img_h) >= UPSCALE_TILE_THRESHOLD: n_tiles = math.ceil(img_w / UPSCALE_TILE_SIZE) * math.ceil(img_h / UPSCALE_TILE_SIZE) logger.info( "Upscale job %s: %dx%d → tiled mode (%d tiles expected). " "scale=%s denoise=%s sharpen=%s timeout=%ds", job_id[:8], img_w, img_h, n_tiles, output_scale, denoise, sharpen, _UPSCALE_TIMEOUT_SECS, ) else: logger.info( "Upscale job %s: %dx%d → whole-image mode. " "scale=%s denoise=%s sharpen=%s", job_id[:8], img_w, img_h, output_scale, denoise, sharpen, ) # --- Optional: denoise BEFORE upscale (on the small original) ---------- if denoise: img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) img_bgr = _denoise(img_bgr) pil_img = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)) # --- Upscale with job-level timeout ------------------------------------ # Run engine.upscale() in a worker thread so we can enforce a hard time # ceiling. If the timeout fires, the job is immediately marked "failed" # so the user gets feedback. The underlying thread becomes a daemon and # will be collected when the server restarts. with ThreadPoolExecutor(max_workers=1) as _pool: _future = _pool.submit(engine.upscale, pil_img) try: sr_img = _future.result(timeout=_UPSCALE_TIMEOUT_SECS) except _FuturesTimeout: elapsed_min = (time.monotonic() - t_job_start) / 60 logger.error( "Upscale job %s TIMED OUT after %.1f min (limit %d min). " "Background thread will continue until computation finishes.", job_id[:8], elapsed_min, _UPSCALE_TIMEOUT_SECS // 60, ) update_job( job_id, status = "failed", error_message = ( f"처리 시간이 너무 오래 걸려 중단되었습니다 " f"({_UPSCALE_TIMEOUT_SECS // 60}분 초과). " f"이미지 크기를 줄이거나 나중에 다시 시도해 주세요." ), ) return # --- Optional: sharpen AFTER upscale, BEFORE downscale ----------------- if sharpen: img_arr = np.array(sr_img) # PIL RGB → numpy RGB uint8 img_arr = _sharpen(img_arr) # unsharp mask (channel-independent) sr_img = Image.fromarray(img_arr) if output_scale == "2x": w, h = sr_img.size sr_img = sr_img.resize((w // 2, h // 2), Image.Resampling.LANCZOS) elif output_scale == "1.5x": # Target: 1.5× the original. Engine always outputs 4×, so the # resize factor is 1.5/4 = 3/8 = 0.375. # round() gives the nearest integer; for the 671×406 test image: # round(2684 × 3/8) = round(1006.5) = 1006 # round(1624 × 3/8) = round(609.0) = 609 # matching 671×1.5≈1006 and 406×1.5=609 as closely as integers allow. w, h = sr_img.size sr_img = sr_img.resize( (round(w * 3 / 8), round(h * 3 / 8)), Image.Resampling.LANCZOS, ) sr_img.save(str(result_path), format="JPEG", quality=95) # Build the URL path relative to the /storage mount. # result_path is an absolute path under STORAGE_DIR, so we compute the # portion after STORAGE_DIR to construct the URL fragment. rel = result_path.relative_to(STORAGE_DIR) elapsed = time.monotonic() - t_job_start logger.info("Upscale job %s completed in %.1fs", job_id[:8], elapsed) update_job(job_id, status="completed", result_image=f"/{rel.as_posix()}") except Exception: update_job( job_id, status="failed", error_message=traceback.format_exc(), ) # --------------------------------------------------------------------------- # POST /upscale/upload # --------------------------------------------------------------------------- @router.post("/upload") async def upscale_upload( file: UploadFile, background_tasks: BackgroundTasks, request: Request, output_scale: str = Form("4x"), denoise: bool = Form(False), sharpen: bool = Form(False), ) -> dict: """Accept an image, EXIF-correct it, persist it, and kick off upscaling. Form fields ----------- file : image file (multipart) output_scale : "4x" (default) | "2x" | "1.5x" denoise : bool (default False) — apply NLM denoising before upscale sharpen : bool (default False) — apply mild unsharp mask after upscale Returns ``{"job_id": }`` immediately; upscaling runs in the background. Poll GET /upscale/result/{job_id} for completion. """ # --- output_scale validation --------------------------------------------- if output_scale not in ("4x", "2x", "1.5x"): raise HTTPException( status_code=422, detail=f"output_scale must be '4x', '2x', or '1.5x', got '{output_scale}'.", ) # --- Content-type guard -------------------------------------------------- if not (file.content_type or "").startswith("image/"): raise HTTPException( status_code=415, detail=f"Expected an image file, got content-type '{file.content_type}'.", ) # --- Read and EXIF-correct ----------------------------------------------- raw_bytes = await file.read() if not raw_bytes: raise HTTPException(status_code=400, detail="Uploaded file is empty.") try: pil_img = Image.open(io.BytesIO(raw_bytes)) pil_img = ImageOps.exif_transpose(pil_img) pil_img = pil_img.convert("RGB") except Exception as exc: raise HTTPException(status_code=422, detail=f"Could not decode image: {exc}") # --- Persist original image ---------------------------------------------- job_id = str(uuid.uuid4()) created_at = datetime.now(timezone.utc) folder = job_folder_name(job_id, created_at) job_dir = _UPSCALE_STORAGE / folder job_dir.mkdir(parents=True, exist_ok=True) original_path = job_dir / "original.jpg" pil_img.save(str(original_path), format="JPEG", quality=95) # --- Create job record --------------------------------------------------- # Pass created_at explicitly (as ISO string, matching UpscaleJob.created_at: str) # so upscale_jobs_store.create_job() doesn't generate a different timestamp. orig_rel = original_path.relative_to(STORAGE_DIR) job = UpscaleJob( id = job_id, original_image = f"/{orig_rel.as_posix()}", status = "processing", output_scale = output_scale, denoise = denoise, sharpen = sharpen, created_at = created_at.isoformat(), ) create_job(job) # --- Enqueue background upscaling ---------------------------------------- engine: BaseUpscaleEngine = request.app.state.upscale_engine background_tasks.add_task( run_upscale_job, job_id = job_id, folder_name = folder, original_path = original_path, engine = engine, output_scale = output_scale, denoise = denoise, sharpen = sharpen, ) return {"job_id": job_id} # --------------------------------------------------------------------------- # GET /upscale/result/{job_id} # --------------------------------------------------------------------------- @router.get("/result/{job_id}") def upscale_result(job_id: str) -> dict: """Return the current status of an upscale job. Response fields: job_id : str status : "pending" | "processing" | "completed" | "failed" result_image_url : str | null — /storage/upscale/{job_id}/result.jpg error_message : str | null """ job = get_job(job_id) if job is None: raise HTTPException(status_code=404, detail=f"Upscale job '{job_id}' not found.") return { "job_id": job.id, "status": job.status, "result_image_url": job.result_image, "error_message": job.error_message, }