app.py

# Import necessary libraries
import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter, ImageOps
import numpy as np
import cv2
import torch
import os
import tempfile
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer

import warnings
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', message='.*torchvision.transforms.functional_tensor.*')

# Set NumExpr threads
import os
os.environ['NUMEXPR_MAX_THREADS'] = '16'

# DeOldify imports
from deoldify.visualize import get_image_colorizer


# Create a temporary directory for saving images
TEMP_DIR = tempfile.gettempdir()

# Global variables for models (loaded once)
face_enhancer = None
bg_upsampler = None
deoldify_colorizer = None

def initialize_deoldify():
    """Initialize DeOldify colorizer (lazy loading)"""
    global deoldify_colorizer
    
    if deoldify_colorizer is None:
        print("Loading DeOldify model...")
        try:
            # Suppress fastai warnings
            import warnings
            warnings.filterwarnings('ignore', message='.*training set is empty.*')
            warnings.filterwarnings('ignore', message='.*validation set is empty.*')
            
            deoldify_colorizer = get_image_colorizer(artistic=False)
            print("✅ DeOldify loaded successfully!")
            return True
        except Exception as e:
            print(f"❌ Error loading DeOldify: {e}")
            return False
    return True

def initialize_face_enhancer():
    """Initialize GFPGAN model (lazy loading)"""
    global face_enhancer, bg_upsampler
    
    if face_enhancer is None:
        print("Loading GFPGAN model...")
        try:
            # Initialize background upsampler (for non-face areas)
            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
            bg_upsampler = RealESRGANer(
                scale=2,
                model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
                model=model,
                tile=400,
                tile_pad=10,
                pre_pad=0,
                half=False
            )
            
            # Initialize GFPGAN
            face_enhancer = GFPGANer(
                model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
                upscale=2,
                arch='clean',
                channel_multiplier=2,
                bg_upsampler=bg_upsampler
            )
            print("✅ GFPGAN loaded successfully!")
            return True
        except Exception as e:
            print(f"❌ Error loading GFPGAN: {e}")
            return False
    return True

def enhance_faces(img, face_enhancement_strength=100):
    """
    Enhance faces in the image using GFPGAN
    
    Parameters:
    - img: PIL Image
    - face_enhancement_strength: 0-100, blend between original and enhanced
    
    Returns:
    - Enhanced PIL Image
    """
    if img is None:
        return img
    
    if not initialize_face_enhancer():
        print("Face enhancer not available, returning original image")
        return img
    
    try:
        # Ensure image is PIL Image
        if not isinstance(img, Image.Image):
            return img
        
        # Convert PIL to numpy array
        img_array = np.asarray(img, dtype=np.uint8)
        
        # Force a copy to ensure it's a proper numpy array
        img_array = np.array(img_array, copy=True, dtype=np.uint8)
        
        # Convert RGB to BGR (manual conversion for compatibility)
        img_cv = img_array[:, :, ::-1].copy()
        
        # Apply GFPGAN
        _, _, output = face_enhancer.enhance(
            img_cv,
            has_aligned=False,
            only_center_face=False,
            paste_back=True,
            weight=face_enhancement_strength / 100
        )
        
        # Convert back to PIL (RGB)
        output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        enhanced_img = Image.fromarray(output_rgb)
        
        return enhanced_img
        
    except Exception as e:
        print(f"Face enhancement error: {e}")
        return img

def remove_scratches(img, strength):
    """Improved scratch and noise removal"""
    if strength == 0:
        return img
    
    if strength < 50:
        kernel_size = 3
    elif strength < 75:
        kernel_size = 5
    else:
        kernel_size = 7
    
    scratch_filtered = img.filter(ImageFilter.MedianFilter(size=kernel_size))
    blur_radius = strength / 50
    noise_filtered = scratch_filtered.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    alpha = min(strength / 100, 1.0) * 0.7
    result = Image.blend(img, noise_filtered, alpha)
    
    return result

# ==================== STEP 1: RESTORATION ====================
def step1_restore(
    image,
    face_enhance,
    face_enhance_strength,
    upscale_factor,
    scratch_removal,
    denoise_level,
    sharpness,
    contrast_mode,
    contrast_manual,
    brightness,
    fade_mode
):
    """Step 1: Basic Restoration"""
    if image is None:
        return None, "❌ Please upload an image first!"
    
    try:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        status_messages = ["🔧 Step 1: Restoration Started"]
        
        # Face Enhancement
        if face_enhance:
            status_messages.append("🎭 Enhancing faces...")
            image = enhance_faces(image, face_enhance_strength)
        
        # Scratch removal
        if scratch_removal > 0:
            status_messages.append("🧹 Removing scratches...")
            image = remove_scratches(image, scratch_removal)
        
        # Denoise
        if denoise_level != "None":
            status_messages.append(f"🎚️ Denoising ({denoise_level})...")
            denoise_map = {"Low": 0.5, "Medium": 1.0, "High": 1.5}
            radius = denoise_map.get(denoise_level, 1.0)
            image = image.filter(ImageFilter.GaussianBlur(radius=radius))
        
        # Contrast adjustment
        if contrast_mode == "Auto":
            status_messages.append("📊 Auto contrast...")
            image = ImageOps.autocontrast(image)
        elif contrast_mode == "Manual":
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(contrast_manual / 50)
        
        # Brightness
        if brightness != 50:
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(brightness / 50)
        
        # Fade correction
        if fade_mode == "Auto":
            status_messages.append("🌈 Correcting fade...")
            image = ImageOps.equalize(image)
        
        # Sharpness
        if sharpness != 50:
            enhancer = ImageEnhance.Sharpness(image)
            image = enhancer.enhance(sharpness / 50)
        
        # Upscaling (skip if face enhancement already upscaled)
        if upscale_factor > 1 and not face_enhance:
            status_messages.append(f"🔍 Upscaling {upscale_factor}x...")
            width, height = image.size
            new_size = (int(width * upscale_factor), int(height * upscale_factor))
            image = image.resize(new_size, Image.Resampling.LANCZOS)
        
        status_messages.append("✅ Step 1 Complete!")
        status_text = "\n".join(status_messages)
        
        return image, status_text
    
    except Exception as e:
        return None, f"❌ Error in Step 1: {str(e)}"

# ==================== STEP 2: COLORIZATION WITH DEOLDIFY ====================
def step2_colorize(restored_img, render_factor):
    """
    Step 2: Colorize using DeOldify
    
    Parameters:
    - restored_img: PIL Image from Step 1
    - render_factor: Quality setting (10-44, higher=better but slower)
    
    Returns:
    - Colorized PIL Image, status message
    """
    if restored_img is None:
        return None, "❌ Please complete Step 1 first!"
    
    if not initialize_deoldify():
        return None, "❌ DeOldify model failed to load"
    
    try:
        status_messages = ["🎨 Step 2: Colorization Started"]
        status_messages.append(f"🖼️ Processing with quality level {render_factor}...")
        
        # Save to temporary file (DeOldify needs file path)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
        restored_img.save(temp_file.name)
        temp_file.close()
        
        # Colorize with DeOldify
        colorized = deoldify_colorizer.get_transformed_image(
            path=temp_file.name,
            render_factor=render_factor,
            watermarked=False
        )
        
        # Clean up temp file
        os.unlink(temp_file.name)
        
        status_messages.append("✅ Step 2 Complete! Realistic colors applied.")
        status_text = "\n".join(status_messages)
        
        return colorized, status_text
        
    except Exception as e:
        return None, f"❌ Error in Step 2: {str(e)}"

# ==================== STEP 3: CREATIVE (Placeholder) ====================
def step3_creative(colorized_img, vintage_film, dreamy_glow, era_style, artistic_filter):
    """Step 3: Creative Effects (Coming Soon)"""
    if colorized_img is None:
        return None, "❌ Please complete Step 2 first!"
    
    # Placeholder - just return the input for now
    status = "🚧 Step 3: Creative effects coming soon!\n✅ Your colorized photo is ready to download."
    return colorized_img, status

# ==================== GRADIO INTERFACE ====================
with gr.Blocks(title="AI Photo Restoration & Colorization", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 📸 AI Photo Restoration & Colorization Studio
    
    Transform your old black & white photos with AI-powered restoration and colorization!
    
    ### 🎯 3-Step Process:
    1. **🔧 Restoration** - Fix damage, enhance faces, improve quality
    2. **🎨 Colorization** - Add realistic colors with DeOldify AI
    3. **✨ Creative** - Apply artistic effects (Coming Soon)
    """)
    
    # Main layout
    with gr.Row():
        # Input column
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="📤 Upload Your Photo",
                type="pil",
                height=350
            )
            gr.Markdown("""
            **📝 Tips:**
            - JPG or PNG format
            - Works best with faces
            - Old B&W photos welcome!
            """)
        
        # State variables (hidden)
        restored_state = gr.State()
        colorized_state = gr.State()
        
        # Output column - Results Gallery
        restored_display = gr.Image(
            label="Restored Photo",
            height=300,
            format="jpg"
        )
        colorized_display = gr.Image(
            label="Colorized Photo",
            height=300,
            format="jpg"
        )
        creative_display = gr.Image(
            label="Creative Photo",
            height=300,
            format="jpg"
        )
    
    gr.Markdown("---")
    
    # ==================== STEP 1: RESTORATION ====================
    with gr.Accordion("🔧 Step 1: Restoration & Enhancement", open=True):
        gr.Markdown("""
        **First, let's restore your photo and fix any damage.**
        Adjust restoration parameters below.
        """)
        
        with gr.Accordion("⚙️ Restoration Settings", open=True):
            with gr.Row():
                face_enhance = gr.Checkbox(
                    label="✨ Face Enhancement (GFPGAN)",
                    value=False,
                    info="AI-powered face restoration"
                )
                
                face_enhance_strength = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=60,
                    step=10,
                    label="Face Enhancement Strength",
                    info="Recommended: 50-80 for natural look"
                )
            
            with gr.Row():
                upscale_factor = gr.Slider(
                    minimum=1,
                    maximum=4,
                    value=2,
                    step=1,
                    label="Upscale Factor",
                    info="Face enhancement includes 2x upscale"
                )
                
                scratch_removal = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=10,
                    step=10,
                    label="Scratch & Dust Removal"
                )
            
            with gr.Row():
                denoise_level = gr.Radio(
                    choices=["None", "Low", "Medium", "High"],
                    value="Low",
                    label="Noise Reduction"
                )
                
                sharpness = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=50,
                    step=10,
                    label="Sharpness (50=neutral)"
                )
            
            with gr.Row():
                contrast_mode = gr.Radio(
                    choices=["None", "Auto", "Manual"],
                    value="Auto",
                    label="Contrast Recovery"
                )
                
                contrast_manual = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=50,
                    step=10,
                    label="Manual Contrast (50=neutral)"
                )
            
            with gr.Row():
                fade_mode = gr.Radio(
                    choices=["None", "Auto"],
                    value="Auto",
                    label="Fade Correction"
                )
                
                brightness = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=50,
                    step=10,
                    label="Brightness (50=neutral)"
                )
        
        step1_status = gr.Textbox(label="📋 Step 1 Status", lines=3, interactive=False)
        
        step1_button = gr.Button("🔧 Restore Photo", variant="primary", size="lg")
    
    # ==================== STEP 2: COLORIZATION ====================
    with gr.Accordion("🎨 Step 2: AI Colorization (DeOldify)", open=False):
        gr.Markdown("""
        **Add realistic colors to your restored photo using DeOldify AI.**
        
        The AI model analyzes your photo and applies natural, historically-accurate colors.
        """)
        
        render_factor = gr.Slider(
                minimum=10,
                maximum=44,
                value=35,
                step=1,
                label="Quality Level",
                info="10=Fast, 35=Balanced, 44=Best"
        )            
        
        colorize_btn = gr.Button("🎨 Colorize Photo", variant="primary", size="lg")
        step2_status = gr.Textbox(label="📋 Step 2 Status", lines=3, interactive=False)
    
    # ==================== STEP 3: CREATIVE ====================
    with gr.Accordion("✨ Step 3: Creative Effects", open=False):
        gr.Markdown("""
        **Finally, add artistic touches and creative effects.**
        Experiment with vintage filters and artistic styles.
        
        🚧 *Coming Soon: These effects are under development!*
        """)
        
        with gr.Accordion("🎨 Creative Settings", open=True):
            with gr.Row():
                vintage_film = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=0,
                    step=5,
                    label="🎞️ Vintage Film Effect",
                    info="Add film grain and nostalgic feel",
                    interactive=False
                )
                
                dreamy_glow = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=0,
                    step=5,
                    label="✨ Dreamy Glow",
                    info="Soft romantic blur effect",
                    interactive=False
                )
            
            with gr.Row():
                era_style = gr.Radio(
                    choices=["None", "1920s Silent Film", "1950s Technicolor", "1970s Polaroid", "Modern HDR"],
                    value="None",
                    label="📅 Era Style Transfer",
                    interactive=False
                )
                
                artistic_filter = gr.Dropdown(
                    choices=["None", "Oil Painting", "Watercolor", "Sketch", "Pop Art"],
                    value="None",
                    label="🎨 Artistic Filter",
                    interactive=False
                )
        
        step3_status = gr.Textbox(label="📋 Step 3 Status", lines=3, interactive=False)
        
        step3_button = gr.Button("✨ Apply Creative Effects", variant="primary", size="lg")
    
    # ==================== BUTTON ACTIONS ====================
    
    # Step 1: Restore
    def step1_process(img, *args):
        result, status = step1_restore(img, *args)
        if result:
            return result, result, status
        return None, None, status
    
    step1_button.click(
        fn=step1_process,
        inputs=[
            input_image, face_enhance, face_enhance_strength, upscale_factor,
            scratch_removal, denoise_level, sharpness,
            contrast_mode, contrast_manual, brightness, fade_mode
        ],
        outputs=[restored_display, restored_state, step1_status]
    )
    
    # Step 2: Colorize with DeOldify
    def step2_process(restored_img, render_factor):
        result, status = step2_colorize(restored_img, render_factor)
        if result:
            return result, result, status
        return None, None, status
    
    colorize_btn.click(
        fn=step2_process,
        inputs=[restored_state, render_factor],
        outputs=[colorized_display, colorized_state, step2_status]
    )
    
    # Step 3: Creative
    def step3_process(colorized_img, *args):
        result, status = step3_creative(colorized_img, *args)
        if result:
            return result, status
        return None, status
    
    step3_button.click(
        fn=step3_process,
        inputs=[colorized_state, vintage_film, dreamy_glow, era_style, artistic_filter],
        outputs=[creative_display, step3_status]
    )
    
    # ==================== FOOTER ====================
    gr.Markdown("""
    ---
    ### 💡 Tips for Best Results:
    
    **Step 1 - Restoration:**
    - Upload your photo in the top-left panel
    - Use Face Enhancement for portraits (50-80% strength for natural look)
    - Apply scratch removal for damaged photos
    - Auto contrast works well for most faded photos
    
    **Step 2 - Colorization:**
    - Quality 35 is recommended for balanced speed/quality
    - Quality 20 for faster preview
    - Quality 42-44 for best results (slower)
    
    **Step 3 - Creative:**
    - Coming soon! Artistic filters and vintage effects
    
    **💾 Downloading Images:**
    - Click the download icon (⬇️) in the corner of any result image to save
    - All images are automatically formatted as high-quality
    
    ### 🔧 Powered by:
    - **GFPGAN v1.3** for AI face restoration
    - **Real-ESRGAN** for background upscaling
    - **DeOldify** for AI colorization
    - **Gradio** for the interface
    """)

if __name__ == "__main__":
    demo.launch()