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