File size: 8,995 Bytes
b2ccdfc 3c55f2b b2ccdfc 3c55f2b b2ccdfc 3c55f2b b2ccdfc 3c55f2b b2ccdfc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 |
"""
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
|