""" Image Enhancement Utility using Real-ESRGAN-ncnn-vulkan Enhances blurry/low-quality images before VLM processing """ import os import subprocess import tempfile import zipfile import shutil from pathlib import Path from PIL import Image import urllib.request import platform class ImageEnhancer: """Handles image enhancement using Real-ESRGAN-ncnn-vulkan""" # Download URLs for Windows executable REALESRGAN_VERSION = "v0.2.0" REALESRGAN_WINDOWS_URL = f"https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan/releases/download/{REALESRGAN_VERSION}/realesrgan-ncnn-vulkan-v0.2.0-windows.zip" def __init__(self, base_dir: str = None): """Initialize image enhancer Args: base_dir: Base directory for storing executable and models """ if base_dir is None: base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) self.base_dir = Path(base_dir) self.enhancer_dir = self.base_dir / "utils" / "realesrgan" self.executable_path = None self.models_path = None self.is_available = False # Initialize enhancer self._setup_enhancer() def _setup_enhancer(self): """Setup Real-ESRGAN enhancer (download if needed)""" try: # Check if already exists if self._check_existing_installation(): print("✅ Real-ESRGAN enhancer already installed") self.is_available = True return # Download and setup print("📥 Downloading Real-ESRGAN enhancer...") self._download_and_extract() if self._check_existing_installation(): print("✅ Real-ESRGAN enhancer installed successfully") self.is_available = True else: print("⚠️ Real-ESRGAN enhancer setup incomplete") self.is_available = False except Exception as e: print(f"⚠️ Failed to setup Real-ESRGAN enhancer: {str(e)}") print(" Image enhancement will be skipped") self.is_available = False def _check_existing_installation(self) -> bool: """Check if Real-ESRGAN is already installed""" if not self.enhancer_dir.exists(): return False # Look for executable exe_name = "realesrgan-ncnn-vulkan.exe" if platform.system() == "Windows" else "realesrgan-ncnn-vulkan" possible_paths = [ self.enhancer_dir / exe_name, self.enhancer_dir / "realesrgan-ncnn-vulkan" / exe_name, self.enhancer_dir / "realesrgan-ncnn-vulkan-v0.2.0-ubuntu" / exe_name, ] for path in possible_paths: if path.exists() and (path.stat().st_mode & 0o111 or platform.system() == "Windows"): self.executable_path = path print(f" Found executable: {path}") # Look for models directory in multiple locations possible_model_dirs = [ path.parent / "models", self.enhancer_dir / "models", ] for models_dir in possible_model_dirs: if models_dir.exists(): self.models_path = models_dir print(f" Found models: {models_dir}") return True # Even without models directory, if executable exists, mark as available # Models can be specified via -n parameter print(f" Models directory not found, but executable is available") self.models_path = path.parent / "models" # Set expected path return True return False def _download_and_extract(self): """Download and extract Real-ESRGAN executable""" if platform.system() != "Windows": print("⚠️ Auto-download only supported on Windows. Please manually install Real-ESRGAN-ncnn-vulkan") return # Create directory self.enhancer_dir.mkdir(parents=True, exist_ok=True) # Download file zip_path = self.enhancer_dir / "realesrgan.zip" print(f" Downloading from {self.REALESRGAN_WINDOWS_URL}...") try: urllib.request.urlretrieve(self.REALESRGAN_WINDOWS_URL, zip_path) except Exception as e: print(f" Download failed: {str(e)}") return # Extract print(" Extracting files...") try: with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(self.enhancer_dir) except Exception as e: print(f" Extraction failed: {str(e)}") return # Cleanup zip file zip_path.unlink() print(" Setup complete!") def enhance_image(self, image_path: str, scale: int = 2, model_name: str = "realesrgan-x4plus") -> str: """Enhance image using Real-ESRGAN Args: image_path: Path to input image scale: Upscale ratio (2, 3, or 4) model_name: Model to use (realesrgan-x4plus, realesrgan-x4plus-anime, realesrnet-x4plus) Returns: Path to enhanced image """ if not self.is_available: print("⚠️ Enhancement not available, using original image") return image_path # Create temporary output file input_path = Path(image_path) output_path = input_path.parent / f"{input_path.stem}_enhanced{input_path.suffix}" try: # Build command cmd = [ str(self.executable_path), "-i", str(image_path), "-o", str(output_path), "-n", model_name, "-s", str(scale), "-f", "png" # Output format ] # Add model path if available and exists if self.models_path and self.models_path.exists(): cmd.extend(["-m", str(self.models_path)]) # Run enhancement result = subprocess.run( cmd, capture_output=True, text=True, timeout=30, creationflags=subprocess.CREATE_NO_WINDOW if platform.system() == "Windows" else 0 ) if result.returncode == 0 and output_path.exists(): print(f"✨ Image enhanced successfully (scale={scale}x)") return str(output_path) else: if result.stderr: print(f"⚠️ Enhancement failed: {result.stderr}") print(" Using original image") return image_path except subprocess.TimeoutExpired: print("⚠️ Enhancement timeout, using original image") return image_path except Exception as e: print(f"⚠️ Enhancement error: {str(e)}, using original image") return image_path def enhance_pil_image(self, pil_image: Image.Image, scale: int = 2, model_name: str = "realesrgan-x4plus") -> Image.Image: """Enhance PIL Image object Args: pil_image: PIL Image object scale: Upscale ratio (2, 3, or 4) model_name: Model to use Returns: Enhanced PIL Image object """ if not self.is_available: return pil_image # Save to temporary file with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_input: temp_input_path = temp_input.name pil_image.save(temp_input_path, "PNG") try: # Enhance enhanced_path = self.enhance_image(temp_input_path, scale, model_name) # Load enhanced image if enhanced_path != temp_input_path: enhanced_image = Image.open(enhanced_path).convert("RGB") # Cleanup enhanced temp file try: os.unlink(enhanced_path) except: pass return enhanced_image else: return pil_image finally: # Cleanup input temp file try: os.unlink(temp_input_path) except: pass # Global enhancer instance _enhancer_instance = None def get_enhancer() -> ImageEnhancer: """Get global enhancer instance""" global _enhancer_instance if _enhancer_instance is None: _enhancer_instance = ImageEnhancer() return _enhancer_instance