app.py

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 """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>AI Face Restoration - Restore Your Photos</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            min-height: 100vh;
            padding: 20px;
        }
        
        .container {
            max-width: 1400px;
            margin: 0 auto;
            background: white;
            border-radius: 20px;
            box-shadow: 0 20px 60px rgba(0,0,0,0.3);
            overflow: hidden;
            padding: 40px;
        }
        
        h1 {
            color: #333;
            margin-bottom: 10px;
            font-size: 2.5em;
        }
        
        .subtitle {
            color: #666;
            margin-bottom: 30px;
            font-size: 1.1em;
        }
        
        .upload-area {
            border: 2px dashed #667eea;
            border-radius: 10px;
            padding: 40px;
            text-align: center;
            background: #f8f9ff;
            transition: all 0.3s ease;
            cursor: pointer;
        }
        
        .upload-area:hover {
            border-color: #764ba2;
            background: #f0f2ff;
        }
        
        .upload-area.drag-over {
            border-color: #764ba2;
            background: #e8ebff;
        }
        
        .file-input {
            display: none;
        }
        
        .upload-icon {
            font-size: 48px;
            margin-bottom: 10px;
        }
        
        .controls {
            margin: 30px 0;
            padding: 20px;
            background: #f8f9ff;
            border-radius: 10px;
        }
        
        .control-group {
            margin-bottom: 15px;
        }
        
        label {
            display: block;
            margin-bottom: 5px;
            color: #555;
            font-weight: 500;
        }
        
        input[type="range"] {
            width: 100%;
            padding: 5px;
        }
        
        .scale-value {
            display: inline-block;
            margin-left: 10px;
            color: #667eea;
            font-weight: bold;
        }
        
        button {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            border: none;
            padding: 12px 30px;
            border-radius: 5px;
            font-size: 16px;
            cursor: pointer;
            transition: transform 0.2s;
            width: 100%;
        }
        
        button:hover {
            transform: translateY(-2px);
        }
        
        button:disabled {
            opacity: 0.6;
            cursor: not-allowed;
        }
        
        .results {
            margin-top: 40px;
        }
        
        .comparison {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 20px;
            margin-top: 20px;
        }
        
        .image-card {
            background: #f8f9ff;
            border-radius: 10px;
            padding: 20px;
            text-align: center;
        }
        
        .image-card h3 {
            margin-bottom: 15px;
            color: #333;
        }
        
        .image-preview {
            max-width: 100%;
            border-radius: 8px;
            box-shadow: 0 4px 8px rgba(0,0,0,0.1);
            cursor: pointer;
            transition: transform 0.3s;
        }
        
        .image-preview:hover {
            transform: scale(1.02);
        }
        
        .download-btn {
            background: #28a745;
            margin-top: 15px;
            padding: 10px;
            font-size: 14px;
        }
        
        .info {
            margin-top: 20px;
            padding: 15px;
            background: #e8ebff;
            border-radius: 8px;
            text-align: center;
        }
        
        .loading {
            text-align: center;
            padding: 40px;
        }
        
        .spinner {
            border: 4px solid #f3f3f3;
            border-top: 4px solid #667eea;
            border-radius: 50%;
            width: 50px;
            height: 50px;
            animation: spin 1s linear infinite;
            margin: 0 auto;
        }
        
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
        
        .error {
            background: #f8d7da;
            color: #721c24;
            padding: 15px;
            border-radius: 8px;
            margin-top: 20px;
        }
        
        .format-badge {
            display: inline-block;
            background: #667eea;
            color: white;
            padding: 2px 8px;
            border-radius: 4px;
            font-size: 12px;
            margin-left: 10px;
        }
        
        @media (max-width: 768px) {
            .container {
                padding: 20px;
            }
            
            .comparison {
                grid-template-columns: 1fr;
            }
            
            h1 {
                font-size: 1.8em;
            }
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>✨ AI Face Restoration</h1>
        <p class="subtitle">
            Restore and enhance faces using GFPGAN v1.4 & RestoreFormer 
            <span class="format-badge">All formats supported</span>
        </p>
        
        <div class="upload-area" id="uploadArea">
            <div class="upload-icon">📸</div>
            <p>Click or drag & drop your image here</p>
            <p style="font-size: 12px; color: #999; margin-top: 10px;">
                Supports: JPG, PNG, WEBP, BMP, TIFF, GIF, HEIC, HEIF
            </p>
            <input type="file" id="fileInput" class="file-input" accept="image/*">
        </div>
        
        <div class="controls">
            <div class="control-group">
                <label>Scale Factor: <span id="scaleValue" class="scale-value">2.0</span></label>
                <input type="range" id="scaleSlider" min="1" max="4" step="0.5" value="2">
            </div>
        </div>
        
        <button id="restoreBtn" disabled>Restore Image 🔄</button>
        
        <div id="results" class="results" style="display: none;"></div>
    </div>
    
    <script>
        const uploadArea = document.getElementById('uploadArea');
        const fileInput = document.getElementById('fileInput');
        const restoreBtn = document.getElementById('restoreBtn');
        const scaleSlider = document.getElementById('scaleSlider');
        const scaleValue = document.getElementById('scaleValue');
        const resultsDiv = document.getElementById('results');
        
        let selectedFile = null;
        let originalPreviewUrl = null;
        
        scaleSlider.addEventListener('input', (e) => {
            scaleValue.textContent = parseFloat(e.target.value).toFixed(1);
        });
        
        uploadArea.addEventListener('click', () => {
            fileInput.click();
        });
        
        uploadArea.addEventListener('dragover', (e) => {
            e.preventDefault();
            uploadArea.classList.add('drag-over');
        });
        
        uploadArea.addEventListener('dragleave', () => {
            uploadArea.classList.remove('drag-over');
        });
        
        uploadArea.addEventListener('drop', (e) => {
            e.preventDefault();
            uploadArea.classList.remove('drag-over');
            const file = e.dataTransfer.files[0];
            if (file && file.type.startsWith('image/')) {
                handleFileSelect(file);
            } else {
                alert('Please drop an image file');
            }
        });
        
        fileInput.addEventListener('change', (e) => {
            if (e.target.files[0]) {
                handleFileSelect(e.target.files[0]);
            }
        });
        
        function handleFileSelect(file) {
            selectedFile = file;
            
            // Create preview
            if (originalPreviewUrl) {
                URL.revokeObjectURL(originalPreviewUrl);
            }
            originalPreviewUrl = URL.createObjectURL(file);
            
            // Show original preview immediately
            resultsDiv.style.display = 'block';
            resultsDiv.innerHTML = `
                <div class="comparison">
                    <div class="image-card">
                        <h3>🖼️ Original Image</h3>
                        <img src="${originalPreviewUrl}" class="image-preview" alt="Original">
                    </div>
                    <div class="image-card">
                        <h3>✨ Restored Preview</h3>
                        <div class="loading">
                            <div class="spinner"></div>
                            <p>Click "Restore Image" to process</p>
                        </div>
                    </div>
                </div>
            `;
            
            restoreBtn.disabled = false;
        }
        
        restoreBtn.addEventListener('click', async () => {
            if (!selectedFile) return;
            
            restoreBtn.disabled = true;
            restoreBtn.textContent = 'Processing... ⏳';
            
            const formData = new FormData();
            formData.append('image', selectedFile);
            formData.append('scale', scaleSlider.value);
            
            try {
                const response = await fetch('/restore', {
                    method: 'POST',
                    body: formData
                });
                
                if (!response.ok) {
                    const error = await response.json();
                    throw new Error(error.detail || 'Restoration failed');
                }
                
                // Get metadata from headers
                const selectedModel = response.headers.get('X-Selected-Model');
                const facesDetected = response.headers.get('X-Faces-Detected');
                const originalFormat = response.headers.get('X-Original-Format');
                const originalFilename = response.headers.get('X-Original-Filename');
                const contentType = response.headers.get('Content-Type');
                
                // Get restored image blob
                const blob = await response.blob();
                const restoredUrl = URL.createObjectURL(blob);
                
                // Extract filename from Content-Disposition or create one
                let filename = `${originalFilename || 'restored'}_restored.${(originalFormat || 'jpg').toLowerCase()}`;
                const contentDisposition = response.headers.get('Content-Disposition');
                if (contentDisposition) {
                    const match = contentDisposition.match(/filename=(.+)/);
                    if (match) filename = match[1];
                }
                
                // Update UI with both images
                resultsDiv.innerHTML = `
                    <div class="comparison">
                        <div class="image-card">
                            <h3>🖼️ Original Image</h3>
                            <img src="${originalPreviewUrl}" class="image-preview" alt="Original">
                            <p style="font-size: 12px; color: #666; margin-top: 10px;">
                                Format: ${originalFormat || 'Unknown'}
                            </p>
                        </div>
                        <div class="image-card">
                            <h3>✨ Restored Image</h3>
                            <img src="${restoredUrl}" class="image-preview" alt="Restored" id="restoredImg">
                            <button class="download-btn" onclick="downloadImage('${restoredUrl}', '${filename}')">
                                💾 Download Restored
                            </button>
                        </div>
                    </div>
                    <div class="info">
                        <strong>🎯 Model Used:</strong> ${selectedModel || 'Auto-selected'} &nbsp;|&nbsp;
                        <strong>👤 Faces Detected:</strong> ${facesDetected || '0'} &nbsp;|&nbsp;
                        <strong>📁 Output Format:</strong> Same as original (${originalFormat || 'JPEG'})
                    </div>
                `;
                
            } catch (error) {
                resultsDiv.innerHTML = `
                    <div class="error">
                        ❌ Error: ${error.message}
                    </div>
                    <div class="comparison">
                        <div class="image-card">
                            <h3>🖼️ Original Image</h3>
                            <img src="${originalPreviewUrl}" class="image-preview" alt="Original">
                        </div>
                    </div>
                `;
            } finally {
                restoreBtn.disabled = false;
                restoreBtn.textContent = 'Restore Image 🔄';
            }
        });
        
        function downloadImage(url, filename) {
            const a = document.createElement('a');
            a.href = url;
            a.download = filename;
            document.body.appendChild(a);
            a.click();
            document.body.removeChild(a);
        }
        
        // Cleanup on page unload
        window.addEventListener('beforeunload', () => {
            if (originalPreviewUrl) {
                URL.revokeObjectURL(originalPreviewUrl);
            }
        });
    </script>
</body>
</html>
    """

# =====================================================
# 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)