import os import sys import uuid import shutil from datetime import datetime import cv2 import torch import numpy as np from PIL import Image import io import base64 from fastapi import FastAPI, File, UploadFile, Form, HTTPException from fastapi.responses import HTMLResponse, FileResponse, JSONResponse, Response from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from torchvision.transforms import functional # Fix for torchvision import issue sys.modules["torchvision.transforms.functional_tensor"] = functional from basicsr.archs.srvgg_arch import SRVGGNetCompact from gfpgan.utils import GFPGANer from realesrgan.utils import RealESRGANer # Try to import HEIC/HEIF support try: import pillow_heif pillow_heif.register_heif_opener() HEIF_SUPPORT = True except ImportError: HEIF_SUPPORT = False print("HEIC/HEIF support not available. Install pillow-heif for full format support.") # ===================================================== # DIRECTORIES # ===================================================== os.makedirs("models", exist_ok=True) os.makedirs("uploads", exist_ok=True) os.makedirs("outputs", exist_ok=True) # ===================================================== # MODEL DOWNLOADS # ===================================================== MODEL_FILES = { "realesr-general-x4v3.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", "GFPGANv1.4.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth", "RestoreFormer.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth" } def download_models(): for filename, url in MODEL_FILES.items(): filepath = os.path.join("models", filename) if not os.path.exists(filepath): print(f"Downloading {filename}...") import urllib.request urllib.request.urlretrieve(url, filepath) print(f"Downloaded {filename}") download_models() # ===================================================== # REAL-ESRGAN # ===================================================== device = "cuda" if torch.cuda.is_available() else "cpu" model = SRVGGNetCompact( num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type="prelu" ) upsampler = RealESRGANer( scale=4, model_path="models/realesr-general-x4v3.pth", model=model, tile=400, tile_pad=20, pre_pad=0, half=torch.cuda.is_available() ) # ===================================================== # PRELOAD FACE MODELS # ===================================================== print("Loading GFPGAN v1.4...") GFPGAN_MODEL = GFPGANer( model_path="models/GFPGANv1.4.pth", upscale=2, arch="clean", channel_multiplier=2, bg_upsampler=upsampler ) print("Loading RestoreFormer...") RESTOREFORMER_MODEL = GFPGANer( model_path="models/RestoreFormer.pth", upscale=2, arch="RestoreFormer", channel_multiplier=2, bg_upsampler=upsampler ) print("Models loaded successfully!") # ===================================================== # HELPER FUNCTIONS # ===================================================== def get_image_format(file_path): """Detect image format from file""" ext = os.path.splitext(file_path)[1].lower().replace('.', '') format_map = { 'jpg': 'JPEG', 'jpeg': 'JPEG', 'png': 'PNG', 'webp': 'WEBP', 'bmp': 'BMP', 'tiff': 'TIFF', 'tif': 'TIFF', 'gif': 'GIF', 'heic': 'HEIC', 'heif': 'HEIF' } return format_map.get(ext, 'JPEG') def save_image_with_format(image, path, format_type): """Save image with specified format""" if format_type.upper() == 'JPEG': cv2.imwrite(path, image, [cv2.IMWRITE_JPEG_QUALITY, 95]) elif format_type.upper() == 'PNG': cv2.imwrite(path, image, [cv2.IMWRITE_PNG_COMPRESSION, 3]) elif format_type.upper() == 'WEBP': cv2.imwrite(path, image, [cv2.IMWRITE_WEBP_QUALITY, 90]) elif format_type.upper() in ['BMP', 'TIFF']: cv2.imwrite(path, image) else: cv2.imwrite(path, image) def read_image_with_format(file_path): """Read image regardless of format (including HEIC/HEIF)""" ext = os.path.splitext(file_path)[1].lower() if ext in ['.heic', '.heif'] and HEIF_SUPPORT: img = Image.open(file_path) img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) return img else: img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED) return img def select_best_model(img): """Auto select best model based on image quality""" h, w = img.shape[:2] gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) blur_score = cv2.Laplacian(gray, cv2.CV_64F).var() contrast = gray.std() # Face detection using OpenCV face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') faces = face_cascade.detectMultiScale(gray, 1.1, 5) # Decision logic if len(faces) == 0: return "GFPGANv1.4" # Better for non-face images if min(h, w) < 512 or blur_score < 80 or contrast < 40: return "RestoreFormer" return "GFPGANv1.4" def detect_faces(image): """Face detection using OpenCV""" gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml') faces = face_cascade.detectMultiScale(gray, 1.1, 5) return faces def restore_image(image_path, scale_factor=2): """Main restoration function""" # Read original image img = read_image_with_format(image_path) if img is None: raise Exception("Invalid or corrupted image file") # Get original format original_format = get_image_format(image_path) # Handle grayscale images if len(img.shape) == 2: img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) h, w = img.shape[:2] # Upscale small images if min(h, w) < 512: scale_factor_upscale = max(512 / h, 512 / w) img = cv2.resize(img, (int(w * scale_factor_upscale), int(h * scale_factor_upscale)), interpolation=cv2.INTER_LANCZOS4) # Detect faces faces = detect_faces(img) print(f"Detected {len(faces)} face(s)") # Select and apply model selected_model = select_best_model(img) if selected_model == "RestoreFormer": face_enhancer = RESTOREFORMER_MODEL else: face_enhancer = GFPGAN_MODEL print(f"Selected Model: {selected_model}") # Enhance image _, _, output = face_enhancer.enhance( img, has_aligned=False, only_center_face=False, paste_back=True, weight=0.5 ) # Apply scaling if needed if scale_factor != 2: h, w = output.shape[:2] interpolation = cv2.INTER_AREA if scale_factor < 2 else cv2.INTER_LANCZOS4 output = cv2.resize(output, (int(w * scale_factor / 2), int(h * scale_factor / 2)), interpolation=interpolation) return output, selected_model, original_format, len(faces) # ===================================================== # CLEANUP FUNCTIONS # ===================================================== def cleanup_old_files(directory, hours=1): """Remove files older than specified hours""" current_time = datetime.now().timestamp() for filename in os.listdir(directory): filepath = os.path.join(directory, filename) if os.path.isfile(filepath): file_age = current_time - os.path.getmtime(filepath) if file_age > hours * 3600: os.remove(filepath) print(f"Cleaned up: {filepath}") # ===================================================== # FASTAPI APP # ===================================================== app = FastAPI( title="AI Face Restoration API", description="Advanced face restoration with GFPGAN v1.4 and RestoreFormer", version="2.0", docs_url=None, # Disable Swagger redoc_url=None # Disable ReDoc ) # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Mount static files app.mount("/outputs", StaticFiles(directory="outputs"), name="outputs") # ===================================================== # API ENDPOINTS # ===================================================== @app.post("/restore") async def restore_endpoint( image: UploadFile = File(...), scale: float = Form(2) ): """Restore uploaded image and return result""" upload_path = None output_filename = None try: # Generate unique filename while preserving extension file_ext = os.path.splitext(image.filename)[1].lower() original_filename = os.path.splitext(image.filename)[0] unique_id = uuid.uuid4().hex[:8] # Save uploaded file upload_path = os.path.join("uploads", f"{unique_id}_{image.filename}") content = await image.read() # Handle HEIC/HEIF specially if file_ext in ['.heic', '.heif'] and HEIF_SUPPORT: heif_file = io.BytesIO(content) pil_image = Image.open(heif_file) pil_image.save(upload_path, format='PNG') else: with open(upload_path, "wb") as f: f.write(content) # Restore image output_img, selected_model, original_format, faces_detected = restore_image(upload_path, scale) # Save restored image with same format output_filename = f"{original_filename}_restored_{unique_id}{file_ext}" output_path = os.path.join("outputs", output_filename) save_image_with_format(output_img, output_path, original_format) # Read output image for preview with open(output_path, "rb") as f: output_bytes = f.read() # Read original for preview with open(upload_path, "rb") as f: original_bytes = f.read() # Return direct image response with metadata return Response( content=output_bytes, media_type=f"image/{original_format.lower()}", headers={ "X-Selected-Model": selected_model, "X-Faces-Detected": str(faces_detected), "X-Original-Format": original_format, "X-Original-Filename": original_filename, "Content-Disposition": f"inline; filename={output_filename}" } ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) finally: # Clean up upload if upload_path and os.path.exists(upload_path): os.remove(upload_path) # Periodic cleanup of old outputs (every 10 requests) if np.random.random() < 0.1: cleanup_old_files("outputs", hours=1) @app.get("/preview/{filename}") async def get_preview(filename: str): """Get preview of restored image""" file_path = os.path.join("outputs", filename) if os.path.exists(file_path): return FileResponse(file_path) raise HTTPException(status_code=404, detail="Preview not found") # ===================================================== # BROWSER UI # ===================================================== @app.get("/", response_class=HTMLResponse) async def browser_ui(): return """ AI Face Restoration - Restore Your Photos

✨ AI Face Restoration

Restore and enhance faces using GFPGAN v1.4 & RestoreFormer All formats supported

📸

Click or drag & drop your image here

Supports: JPG, PNG, WEBP, BMP, TIFF, GIF, HEIC, HEIF

""" # ===================================================== # HEALTH CHECK (for Hugging Face) # ===================================================== @app.get("/health") async def health_check(): return {"status": "healthy", "models_loaded": True} # ===================================================== # CLEANUP ON SHUTDOWN # ===================================================== @app.on_event("shutdown") async def shutdown_event(): """Cleanup outputs on shutdown""" cleanup_old_files("outputs", hours=0) cleanup_old_files("uploads", hours=0) # ===================================================== # RUN SERVER # ===================================================== if __name__ == "__main__": import uvicorn port = int(os.environ.get("PORT", 7860)) uvicorn.run(app, host="0.0.0.0", port=port)