app.py

# app.py
# Image Upscale and Enhancement with Multiple Models
# By FebryEnsz
# SDK: Gradio
# Hosted on Hugging Face Spaces

import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageEnhance
import cv2
import os
import sys
import subprocess
import time
from huggingface_hub import hf_hub_download

# Create cache directory for models
CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "image_enhancer")
os.makedirs(CACHE_DIR, exist_ok=True)

# Set up logging
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Install required packages at runtime for Hugging Face Spaces
def install_dependencies():
    logger.info("Checking and installing dependencies...")

    packages_to_install = [
        "opencv-python",
        "opencv-contrib-python",  # For dnn_superres module
        "numpy",
        "pillow",
        "torch torchvision torchaudio", # Let pip handle the specific wheels
        "facexlib", # Dependency for GFPGAN
        "basicsr",  # Dependency for RealESRGAN/GFPGAN
        "gfpgan",
        "realesrgan",
        "huggingface_hub" # Ensure hf_hub_download is available
    ]

    # Use a standard index-url or let pip find the best one
    # Forcing CPU might prevent GPU usage if available
    # Let's try without forcing CPU first, Hugging Face Spaces often handles this.
    # If you specifically need CPU only, you might re-add --index-url https://download.pytorch.org/whl/cpu
    
    for package in packages_to_install:
        try:
            logger.info(f"Installing {package}")
            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
        except Exception as e:
            logger.warning(f"Error installing {package}: {str(e)}")

    logger.info("Dependencies installation complete")

# Try to install dependencies on startup
try:
    install_dependencies()
    # Import libraries AFTER installation
    import cv2
    import torch
    import numpy as np
    from PIL import Image, ImageEnhance
    from huggingface_hub import hf_hub_download
    try:
        from realesrgan import RealESRGAN
    except ImportError:
        logger.warning("RealESRGAN import failed after installation attempt.")
        RealESRGAN = None # Set to None if import fails
    try:
        from gfpgan import GFPGANer
    except ImportError:
         logger.warning("GFPGANer import failed after installation attempt.")
         GFPGANer = None # Set to None if import fails

    time.sleep(2)  # Give some time for packages to settle
except Exception as e:
    logger.error(f"Failed to install dependencies or import libraries: {str(e)}")

# Check for GPU or CPU AFTER torch is potentially installed
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")

# Dictionary of available models and their configuration
MODEL_OPTIONS = {
    "OpenCV Super Resolution": {
        "type": "upscale",
        "method": "opencv",
        "scale": 4
    },
    "Real-ESRGAN-x4": {
        "repo_id": "xinntao/Real-ESRGAN",
        "filename": "RealESRGAN_x4plus.pth",
        "type": "upscale",
        "method": "realesrgan",
        "scale": 4
    },
    "GFPGAN (Face Enhancement)": {
        "repo_id": "TencentARC/GFPGAN",
        "filename": "GFPGANv1.4.pth",
        "type": "face",
        "method": "gfpgan",
        "scale": 1 # GFPGAN is primarily for face restoration, upscaling is secondary/handled by bg_upsampler
    },
    "HDR Enhancement": {
        "type": "hdr",
        "method": "custom",
        "scale": 1
    }
}

# Cache for loaded models
model_cache = {}

# Function to load the selected model with robust fallbacks
def load_model(model_name):
    global model_cache

    # Return cached model if available
    if model_name in model_cache:
        logger.info(f"Using cached model: {model_name}")
        return model_cache[model_name]

    logger.info(f"Loading model: {model_name}")
    config = MODEL_OPTIONS.get(model_name)
    if not config:
        return None, f"Model {model_name} not found in configuration"

    model_type = config["type"]

    try:
        # OpenCV based models (always available as fallback if opencv-contrib is installed)
        if config["method"] == "opencv":
            logger.info("Loading OpenCV Super Resolution model")
            try:
                sr = cv2.dnn_superres.DnnSuperResImpl_create()

                # Use EDSR as default model
                model_path = hf_hub_download(
                    repo_id="eugenesiow/edsr",
                    filename="EDSR_x4.pb",
                    cache_dir=CACHE_DIR
                )

                sr.readModel(model_path)
                sr.setModel("edsr", 4)

                # Set backend to cuda if available
                if torch.cuda.is_available():
                    sr.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
                    sr.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

                model_cache[model_name] = (sr, model_type)
                return sr, model_type
            except Exception as e:
                 logger.error(f"Error loading OpenCV SR model: {str(e)}")
                 # Fallback to None if OpenCV SR fails
                 return None, f"Failed to load OpenCV SR model: {str(e)}"


        # Real-ESRGAN models
        elif config["method"] == "realesrgan":
            if RealESRGAN is None:
                 logger.warning("RealESRGAN class not found, falling back to OpenCV SR.")
                 return load_model("OpenCV Super Resolution") # Fallback

            try:
                logger.info("Loading Real-ESRGAN model")

                model_path = hf_hub_download(
                    repo_id=config["repo_id"],
                    filename=config["filename"],
                    cache_dir=CACHE_DIR
                )

                # Initialize RealESRGAN with the correct device
                model = RealESRGAN(device, scale=config["scale"])
                model.load_weights(model_path)

                model_cache[model_name] = (model, model_type)
                return model, model_type
            except Exception as e:
                logger.error(f"Error loading Real-ESRGAN model: {str(e)}")
                logger.warning("Falling back to OpenCV Super Resolution")
                return load_model("OpenCV Super Resolution") # Fallback

        # GFPGAN for face enhancement
        elif config["method"] == "gfpgan":
            if GFPGANer is None:
                 logger.warning("GFPGANer class not found, falling back to OpenCV SR.")
                 return load_model("OpenCV Super Resolution") # Fallback

            try:
                logger.info("Loading GFPGAN model")

                model_path = hf_hub_download(
                    repo_id=config["repo_id"],
                    filename=config["filename"],
                    cache_dir=CACHE_DIR
                )

                # GFPGANer initialization
                # Note: If you want background upsampling with GFPGAN, you need to initialize bg_upsampler
                # e.g., bg_upsampler=RealESRGANer(model_path='...', model_name='RealESRGAN_x4plus.pth', ...)
                # For simplicity and focusing on face, bg_upsampler=None is used here.
                face_enhancer = GFPGANer(
                    model_path=model_path,
                    upscale=config["scale"], # This upscale might be ignored if paste_back is True and no bg_upsampler
                    arch='clean', # Use 'clean' arch for GFPGANv1.4
                    channel_multiplier=2,
                    bg_upsampler=None # No background upsampling
                )

                model_cache[model_name] = (face_enhancer, model_type)
                return face_enhancer, model_type
            except Exception as e:
                logger.error(f"Error loading GFPGAN model: {str(e)}")
                logger.warning("Falling back to OpenCV Super Resolution")
                return load_model("OpenCV Super Resolution") # Fallback

        # HDR Enhancement (custom implementation)
        elif config["method"] == "custom":
            # No model to load for custom HDR
            model_cache[model_name] = (None, model_type)
            return None, model_type

        else:
            return None, f"Unknown model method: {config['method']}"

    except Exception as e:
        logger.error(f"Unexpected error during model loading for {model_name}: {str(e)}")
        import traceback
        traceback.print_exc()

        # Always provide a fallback method if the desired one completely fails
        if model_name != "OpenCV Super Resolution":
            logger.info("Critical error loading model, falling back to OpenCV Super Resolution")
            return load_model("OpenCV Super Resolution")
        else:
            # If OpenCV SR itself fails, something is fundamentally wrong
            return None, f"Failed to load any model, including fallback: {str(e)}"


# Function to preprocess image for processing
def preprocess_image(image):
    """Convert PIL image to numpy array for processing"""
    if image is None:
        return None

    if isinstance(image, Image.Image):
        # Convert PIL image to numpy array
        img = np.array(image)
    else:
        # Assume it's already a numpy array (e.g., from Gradio internal handling)
        img = image

    # Handle grayscale images by converting to RGB
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

    # Handle RGBA images by removing alpha channel
    if img.shape[2] == 4:
        img = img[:, :, :3]

    # Convert RGB to BGR for OpenCV processing
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    return img_bgr

# Function to postprocess image for display
def postprocess_image(img_bgr):
    """Convert processed BGR image back to RGB PIL image"""
    if img_bgr is None:
        return None

    # Ensure image is uint8
    if img_bgr.dtype != np.uint8:
        # Ensure the range is correct before casting
        img_bgr = np.clip(img_bgr, 0, 255)
        img_bgr = img_bgr.astype(np.uint8)

    # Convert BGR to RGB for PIL
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    return Image.fromarray(img_rgb)

# HDR enhancement function
def enhance_hdr(img_bgr, strength=1.0):
    """Custom HDR enhancement using OpenCV"""
    # Convert BGR to RGB
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    # Convert to float32 for processing, range [0, 1]
    img_float = img_rgb.astype(np.float32) / 255.0

    # --- Exposure Fusion based approach (more robust) ---
    try:
        # Estimate camera response function (merge_mertens is more robust)
        merge_mertens = cv2.createMergeMertens(contrast_weight=1.0, saturation_weight=1.0, exposure_weight=0.0)
        # You'd ideally need multiple exposures for true HDR merge.
        # Simulating this by generating slightly adjusted exposures might not be ideal.
        # Let's use a simpler single-image tone mapping or CLAHE on different channels.
        
        # Using CLAHE on L channel (from LAB) and potentially V channel (from HSV)
        img_lab = cv2.cvtColor(img_float, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(img_lab)

        # Apply CLAHE to L channel
        # ClipLimit proportional to strength
        clahe_l = cv2.createCLAHE(clipLimit=max(1.0, 5.0 * strength), tileGridSize=(8, 8))
        # CLAHE works on uint8, so scale L channel
        l_uint8 = np.clip(l * 255.0, 0, 255).astype(np.uint8)
        l_enhanced_uint8 = clahe_l.apply(l_uint8)
        l_enhanced = l_enhanced_uint8.astype(np.float32) / 255.0

        # Blend original and enhanced L channel based on strength
        l_final = l * (1 - strength) + l_enhanced * strength

        # Merge LAB and convert back to RGB
        img_lab_enhanced = cv2.merge([l_final, a, b])
        img_rgb_enhanced = cv2.cvtColor(img_lab_enhanced, cv2.COLOR_LAB2RGB)

        # --- Additional Enhancements (optional, based on strength) ---
        # Vibrance/Saturation adjustment (HSV)
        img_hsv = cv2.cvtColor(img_rgb_enhanced, cv2.COLOR_RGB2HSV)
        h, s, v = cv2.split(img_hsv)

        # Increase saturation, more for less saturated pixels
        saturation_factor = 0.4 * strength # Adjust factor as needed
        s_enhanced = np.clip(s + (s * saturation_factor * (1 - s)), 0, 1)

        # Slight brightness adjustment
        brightness_factor = 0.1 * strength
        v_enhanced = np.clip(v + (v * brightness_factor), 0, 1)


        # Merge HSV and convert back to RGB
        img_rgb_enhanced_hsv = cv2.cvtColor(cv2.merge([h, s_enhanced, v_enhanced]), cv2.COLOR_HSV2RGB)

        # --- Subtle Detail Enhancement (Unsharp Masking effect) ---
        # Convert back to uint8 for blurring
        img_uint8_detail = (np.clip(img_rgb_enhanced_hsv, 0, 1) * 255).astype(np.uint8)
        blur = cv2.GaussianBlur(img_uint8_detail, (0, 0), 5) # Kernel size 5, sigma automatically calculated
        # Convert blur back to float for calculation
        blur_float = blur.astype(np.float32) / 255.0

        detail = img_rgb_enhanced_hsv - blur_float
        # Add detail back, scaled by strength
        img_final_float = np.clip(img_rgb_enhanced_hsv + detail * (0.8 * strength), 0, 1)

        # Convert back to BGR (uint8) for output
        img_bgr_enhanced = (img_final_float * 255).astype(np.uint8)
        img_bgr_enhanced = cv2.cvtColor(img_bgr_enhanced, cv2.COLOR_RGB2BGR)

        return img_bgr_enhanced

    except Exception as e:
        logger.error(f"Error during HDR enhancement: {str(e)}")
        # Return original image if enhancement fails
        return img_bgr


# Main image enhancement function
def enhance_image(image, model_name, strength=1.0, denoise=0.0, sharpen=0.0):
    """Enhance image using selected model with additional processing options"""
    if image is None:
        return "Please upload an image.", None

    try:
        # Load model
        model, model_info = load_model(model_name)
        if isinstance(model_info, str) and model_info.startswith("Failed"):
             # If loading fails, model is None, info is the error message
             return model_info, None

        model_type = model_info # model_info now holds the model type string

        # Preprocess image
        img_bgr = preprocess_image(image)
        if img_bgr is None:
            return "Failed to process image", None

        # Apply denoising if requested
        if denoise > 0:
            logger.info(f"Applying denoising with strength {denoise}")
            # Adjust h and hColor based on denoise slider
            # Recommended range for h is 10 for color images (adjust based on noise level)
            h_val = int(denoise * 20 + 10) # Map 0-1 slider to approx 10-30 h value
            img_bgr = cv2.fastNlMeansDenoisingColored(
                img_bgr, None,
                h=h_val,
                hColor=h_val,
                templateWindowSize=7,
                searchWindowSize=21
            )

        output_bgr = img_bgr # Initialize output with potentially denoised image

        # Process based on model type
        if model_type == "upscale":
            if model is None:
                 return f"Upscaling model '{model_name}' is not loaded or available.", None
            logger.info(f"Upscaling image with {model_name}")
            
            if model_name == "OpenCV Super Resolution":
                # OpenCV super resolution
                output_bgr = model.upsample(img_bgr)

            elif model_name == "Real-ESRGAN-x4":
                # Real-ESRGAN upscaling
                # Real-ESRGAN model object has a 'predict' method
                output_bgr = model.predict(img_bgr)

            # No else needed, as load_model should handle fallbacks

        elif model_type == "face":
            if model is None:
                 return f"Face enhancement model '{model_name}' is not loaded or available.", None
            logger.info(f"Enhancing face with {model_name}")

            if model_name == "GFPGAN (Face Enhancement)":
                # GFPGAN model object has an 'enhance' method
                try:
                    # GFPGAN returns (cropped_faces, restored_faces, restored_img)
                    # restored_img is the pasted-back result
                    _, _, output_bgr = model.enhance(
                        img_bgr,
                        has_aligned=False,
                        only_center_face=False,
                        paste_back=True
                    )
                except Exception as e:
                    logger.error(f"Error enhancing face with GFPGAN: {str(e)}")
                    # If GFPGAN fails, don't just return, try basic upscaling or original
                    # For now, let's just log and return original or denoised image
                    output_bgr = img_bgr # Keep the denoised (or original) image
                    return f"Error applying GFPGAN: {str(e)}. Returning base image.", postprocess_image(output_bgr)

        elif model_type == "hdr":
            # HDR enhancement doesn't use an external model object, it's a function call
            logger.info(f"Applying HDR enhancement with strength {strength}")
            output_bgr = enhance_hdr(img_bgr, strength=strength)

        else:
            # Should not happen if MODEL_OPTIONS is correct
            return f"Unknown model type for processing: {model_type}", None


        # Apply sharpening if requested (apply to the output of the main process)
        if sharpen > 0:
            logger.info(f"Applying sharpening with strength {sharpen}")
            # Simple unsharp mask effect
            kernel = np.array([
                [0, -1, 0],
                [-1, 5, -1],
                [0, -1, 0]
            ], np.float32)
            # We can adjust the strength by blending original and sharpened, or using a kernel with varying center weight
            # A simpler approach is blending:
            sharpened_img = cv2.filter2D(output_bgr, -1, kernel)
            # Blend original output and sharpened output
            output_bgr = cv2.addWeighted(output_bgr, 1.0 - sharpen, sharpened_img, sharpen, 0)


        # Post-process and return image
        enhanced_image = postprocess_image(output_bgr)

        return "Image enhanced successfully!", enhanced_image

    except Exception as e:
        logger.error(f"An error occurred during image processing: {str(e)}")
        import traceback
        traceback.print_exc()
        # Attempt to return original image on error
        if image is not None:
             try:
                 original_img_pil = Image.fromarray(cv2.cvtColor(preprocess_image(image), cv2.COLOR_BGR2RGB))
                 return f"Processing failed: {str(e)}. Returning original image.", original_img_pil
             except Exception as post_e:
                 logger.error(f"Failed to return original image after error: {str(post_e)}")
                 return f"Processing failed: {str(e)}. Could not return image.", None
        else:
             return f"Processing failed: {str(e)}. No image provided.", None


# Gradio interface
with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
    gr.Markdown(
        """
        # 🖼️ Image Upscale & Enhancement
        ### By FebryEnsz

        Upload an image and enhance it with AI-powered upscaling and enhancement.

        **Features:**
        - Super-resolution upscaling (4x) using Real-ESRGAN or OpenCV
        - Face enhancement for portraits using GFPGAN
        - HDR enhancement for better contrast and details
        - Additional Denoise and Sharpen options
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="Upload Image", type="pil", image_mode="RGB") # Explicitly request RGB
            
            # Changed gr.Box() to gr.Group()
            with gr.Group(): # Replaced gr.Box()
                gr.Markdown("### Enhancement Options")
                model_choice = gr.Dropdown(
                    choices=list(MODEL_OPTIONS.keys()),
                    label="Model Selection",
                    value="OpenCV Super Resolution",
                    allow_flagging="never" # Optional: disable flagging
                )

                with gr.Accordion("Advanced Settings", open=False):
                    # Keep strength_slider visible but update label based on model
                    strength_slider = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        step=0.05, # Added more steps for finer control
                        label="Enhancement Strength", # Default label
                        value=0.8,
                        visible=True # Ensure it's visible
                    )

                    denoise_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05, # Added more steps
                        label="Noise Reduction Strength",
                        value=0.0,
                    )

                    sharpen_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05, # Added more steps
                        label="Sharpening Strength",
                        value=0.0,
                    )

            enhance_button = gr.Button("✨ Enhance Image", variant="primary")

        with gr.Column(scale=1):
            output_text = gr.Textbox(label="Status")
            output_image = gr.Image(label="Enhanced Image", type="pil") # Specify type="pil" consistently

    # Handle model change to update UI
    # This function only needs to update the label of the strength slider
    def on_model_change(model_name):
        model_config = MODEL_OPTIONS.get(model_name, {})
        model_type = model_config.get("type", "")

        if model_type == "hdr":
            return gr.update(label="HDR Intensity")
        elif model_type == "face":
             return gr.update(label="Face Enhancement Strength")
        elif model_type == "upscale":
             return gr.update(label="Enhancement Strength") # Keep a generic label for upscale
        else:
             return gr.update(label="Enhancement Strength") # Default

    model_choice.change(on_model_change, inputs=[model_choice], outputs=[strength_slider])

    # Connect button to function
    enhance_button.click(
        fn=enhance_image,
        inputs=[image_input, model_choice, strength_slider, denoise_slider, sharpen_slider],
        outputs=[output_text, output_image],
        api_name="enhance" # Optional: give it an API name
    )

    # Footer information
    gr.Markdown(
        """
        ### Tips
        - For best results with face enhancement, ensure faces are clearly visible.
        - HDR enhancement works best with images that have both bright and dark areas.
        - For noisy images, try increasing the noise reduction slider.
        - Sharpening can add detail but may also increase noise if applied too strongly.

        ---
        Version 2.1 | Running on: """ + (f"GPU 🚀 ({torch.cuda.get_device_name(0)})" if torch.cuda.is_available() else "CPU ⚙️") + """
        """
    )

# Launch the app
if __name__ == "__main__":
    # Use share=True for a temporary public link (useful for debugging, but not needed for Spaces)
    # Use enable_queue=True for better handling of concurrent requests on Spaces
    demo.launch(enable_queue=True)