From 5374eba08fa76f5deebcb4271a293c901d42a89e Mon Sep 17 00:00:00 2001 From: LemonNexus Date: Wed, 11 Feb 2026 17:32:40 +0000 Subject: [PATCH] feat: add Real-ESRGAN AI upscaling support - New module: modules/ai_upscaler.py - Real-ESRGAN integration for AI-powered upscaling - Designed specifically for low-res game textures/icons - 4x upscale factor with neural network enhancement - Falls back gracefully if model not available - Updated TGA Converter UI: - New "AI Enhanced (Real-ESRGAN)" option in upscale dropdown - Only shown if Real-ESRGAN is installed - AI upscaling happens before canvas placement - Shows progress during AI model loading Real-ESRGAN is much better than basic methods for rendered (non-pixel) icons! To use AI upscaling: 1. pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu 2. pip install realesrgan 3. Download RealESRGAN_x4plus.pth model 4. Select "AI Enhanced" in the UI --- modules/ai_upscaler.py | 220 +++++++++++++++++++++++++++++++++++++ ui/tga_converter_dialog.py | 73 ++++++++++-- 2 files changed, 286 insertions(+), 7 deletions(-) create mode 100644 modules/ai_upscaler.py diff --git a/modules/ai_upscaler.py b/modules/ai_upscaler.py new file mode 100644 index 0000000..3f155d0 --- /dev/null +++ b/modules/ai_upscaler.py @@ -0,0 +1,220 @@ +""" +Lemontropia Suite - AI Image Upscaling with Real-ESRGAN +Optional AI-powered upscaling for game icons and textures. + +Real-ESRGAN is specifically designed for low-resolution game graphics +and produces excellent results for rendered icons (not pixel art). +""" + +import logging +from pathlib import Path +from typing import Optional, Tuple +import numpy as np + +try: + from PIL import Image + PIL_AVAILABLE = True +except ImportError: + PIL_AVAILABLE = False + +logger = logging.getLogger(__name__) + +# Try to import Real-ESRGAN +try: + import torch + from basicsr.archs.rrdbnet_arch import RRDBNet + from realesrgan import RealESRGANer + REALESRGAN_AVAILABLE = True +except ImportError: + REALESRGAN_AVAILABLE = False + logger.info("Real-ESRGAN not available. Install with: pip install realesrgan") + + +class AIIconUpscaler: + """ + AI-powered upscaler for game icons using Real-ESRGAN. + + Real-ESRGAN is trained specifically on game textures and produces + excellent results for low-resolution rendered graphics (not pixel art). + + Usage: + upscaler = AIIconUpscaler() + if upscaler.is_available(): + result = upscaler.upscale(image, scale=4) + """ + + # Model download URLs + MODEL_URLS = { + 'RealESRGAN_x4plus': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth', + 'RealESRGAN_x4plus_anime_6B': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus_anime_6B.pth', + 'RealESRGAN_x2plus': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x2plus.pth', + } + + def __init__(self, model_name: str = 'RealESRGAN_x4plus', device: str = 'cpu'): + """ + Initialize AI upscaler. + + Args: + model_name: Which model to use + device: 'cpu' or 'cuda' (GPU) + """ + self.model_name = model_name + self.device = device + self.upsampler = None + + if REALESRGAN_AVAILABLE: + self._init_model() + + def _init_model(self): + """Initialize the Real-ESRGAN model.""" + try: + # Model parameters + if 'anime' in self.model_name: + # Anime model (6B parameters) + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, + num_block=6, num_grow_ch=32, scale=4) + elif 'x2' in self.model_name: + # 2x upscale + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, + num_block=23, num_grow_ch=32, scale=2) + else: + # 4x upscale (default) + model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, + num_block=23, num_grow_ch=32, scale=4) + + # Get model path + model_path = self._get_model_path() + + if not model_path.exists(): + logger.warning(f"Model not found: {model_path}") + logger.info(f"Download from: {self.MODEL_URLS.get(self.model_name)}") + return + + # Initialize upsampler + self.upsampler = RealESRGANer( + scale=4 if 'x4' in self.model_name else 2, + model_path=str(model_path), + model=model, + tile=0, # No tiling (process whole image) + tile_pad=10, + pre_pad=0, + half=False if self.device == 'cpu' else True, + device=torch.device(self.device) + ) + + logger.info(f"Real-ESRGAN initialized: {self.model_name} on {self.device}") + + except Exception as e: + logger.error(f"Failed to initialize Real-ESRGAN: {e}") + self.upsampler = None + + def _get_model_path(self) -> Path: + """Get path to model file.""" + # Store models in user's home directory + model_dir = Path.home() / ".lemontropia" / "models" + model_dir.mkdir(parents=True, exist_ok=True) + return model_dir / f"{self.model_name}.pth" + + def is_available(self) -> bool: + """Check if AI upscaler is available and ready.""" + return REALESRGAN_AVAILABLE and self.upsampler is not None + + def upscale(self, image: Image.Image, scale: int = 4) -> Optional[Image.Image]: + """ + Upscale an image using Real-ESRGAN. + + Args: + image: PIL Image to upscale + scale: Upscale factor (2 or 4) + + Returns: + Upscaled PIL Image or None if failed + """ + if not self.is_available(): + logger.error("AI upscaler not available") + return None + + try: + # Convert PIL to numpy + img_np = np.array(image) + + # Remove alpha channel if present (Real-ESRGAN expects RGB) + has_alpha = img_np.shape[-1] == 4 + if has_alpha: + alpha = img_np[:, :, 3] + img_rgb = img_np[:, :, :3] + else: + img_rgb = img_np + + # Upscale + output, _ = self.upsampler.enhance(img_rgb, outscale=scale) + + # Restore alpha channel if needed + if has_alpha: + # Upscale alpha with simple resize + from PIL import Image as PILImage + alpha_pil = PILImage.fromarray(alpha) + alpha_upscaled = alpha_pil.resize( + (output.shape[1], output.shape[0]), + PILImage.Resampling.LANCZOS + ) + alpha_np = np.array(alpha_upscaled) + output = np.dstack([output, alpha_np]) + + # Convert back to PIL + result = Image.fromarray(output) + return result + + except Exception as e: + logger.error(f"Upscaling failed: {e}") + return None + + def get_info(self) -> dict: + """Get information about the upscaler status.""" + return { + 'available': self.is_available(), + 'model': self.model_name, + 'device': self.device, + 'model_path': str(self._get_model_path()), + 'dependencies': { + 'torch': torch.__version__ if REALESRGAN_AVAILABLE else 'not installed', + 'realesrgan': REALESRGAN_AVAILABLE + } + } + + +# Convenience function +def upscale_with_ai(image: Image.Image, scale: int = 4) -> Optional[Image.Image]: + """Quick AI upscale using Real-ESRGAN.""" + upscaler = AIIconUpscaler() + return upscaler.upscale(image, scale) + + +def check_ai_upscaler(): + """Check if AI upscaler is available and print status.""" + upscaler = AIIconUpscaler() + info = upscaler.get_info() + + print("=" * 50) + print("AI Upscaler Status (Real-ESRGAN)") + print("=" * 50) + print(f"Available: {info['available']}") + print(f"Model: {info['model']}") + print(f"Device: {info['device']}") + print(f"Model Path: {info['model_path']}") + print(f"PyTorch: {info['dependencies']['torch']}") + print("=" * 50) + + if not info['available']: + print("\nTo install AI upscaler:") + print("1. pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu") + print("2. pip install realesrgan") + print("3. Download model from:") + for name, url in AIIconUpscaler.MODEL_URLS.items(): + print(f" {name}: {url}") + + return info['available'] + + +if __name__ == "__main__": + check_ai_upscaler() \ No newline at end of file diff --git a/ui/tga_converter_dialog.py b/ui/tga_converter_dialog.py index 064282e..567d333 100644 --- a/ui/tga_converter_dialog.py +++ b/ui/tga_converter_dialog.py @@ -17,6 +17,7 @@ from PyQt6.QtCore import Qt, QThread, pyqtSignal from PyQt6.QtGui import QPixmap from modules.tga_converter import TGAConverter +from modules.ai_upscaler import AIIconUpscaler, REALESRGAN_AVAILABLE logger = logging.getLogger(__name__) @@ -59,6 +60,15 @@ class TGAConvertWorker(QThread): total = len(tga_files) success = 0 + # Initialize AI upscaler if needed + ai_upscaler = None + if self.upscale_method == 'ai' and REALESRGAN_AVAILABLE: + self.progress_update.emit("Loading AI upscaler (Real-ESRGAN)...") + ai_upscaler = AIIconUpscaler(device='cpu') + if not ai_upscaler.is_available(): + self.progress_update.emit("AI model not found, falling back to HQ4x") + self.upscale_method = 'hq4x' + canvas_info = f" ({self.canvas_size[0]}x{self.canvas_size[1]} canvas)" if self.canvas_size else "" self.progress_update.emit(f"Found {total} TGA files to convert{canvas_info}") @@ -68,12 +78,17 @@ class TGAConvertWorker(QThread): self.progress_update.emit(f"[{i+1}/{total}] Converting: {tga_path.name}") - output_path = self.converter.convert_tga_to_png( - tga_path, - canvas_size=self.canvas_size, - upscale=self.upscale, - upscale_method=self.upscale_method - ) + # Handle AI upscaling separately + if self.upscale_method == 'ai' and ai_upscaler and ai_upscaler.is_available(): + output_path = self._convert_with_ai(tga_path, ai_upscaler) + else: + output_path = self.converter.convert_tga_to_png( + tga_path, + canvas_size=self.canvas_size, + upscale=True, + upscale_method=self.upscale_method + ) + if output_path: success += 1 self.file_converted.emit(tga_path.name, str(output_path)) @@ -83,6 +98,44 @@ class TGAConvertWorker(QThread): except Exception as e: self.conversion_error.emit(str(e)) + def _convert_with_ai(self, tga_path: Path, ai_upscaler: AIIconUpscaler) -> Optional[Path]: + """Convert a single TGA file using AI upscaling.""" + try: + from PIL import Image + + # Load TGA + image = Image.open(tga_path) + if image.mode != 'RGBA': + image = image.convert('RGBA') + + # AI upscale (4x) + self.progress_update.emit(f" AI upscaling {tga_path.name}...") + upscaled = ai_upscaler.upscale(image, scale=4) + + if upscaled is None: + return None + + # Apply canvas if requested + if self.canvas_size: + upscaled = self.converter._apply_canvas( + upscaled, self.canvas_size, upscale=False + ) + + # Save + output_path = self.converter.output_dir / f"{tga_path.stem}.png" + upscaled.save(output_path, 'PNG') + + return output_path + + except Exception as e: + logger.error(f"AI conversion failed: {e}") + return None + + self.conversion_complete.emit(success, total) + + except Exception as e: + self.conversion_error.emit(str(e)) + def stop(self): """Stop the conversion.""" self._is_running = False @@ -204,10 +257,16 @@ class TGAConverterDialog(QDialog): self.upscale_method_combo.addItem("Sharp Pixels (NEAREST)", "nearest") self.upscale_method_combo.addItem("Smooth (HQ4x-style)", "hq4x") self.upscale_method_combo.addItem("Photorealistic (LANCZOS)", "lanczos") + + # Add AI option if Real-ESRGAN is available + if REALESRGAN_AVAILABLE: + self.upscale_method_combo.addItem("🤖 AI Enhanced (Real-ESRGAN)", "ai") + self.upscale_method_combo.setToolTip( "NEAREST: Sharp pixel edges (best for pixel art)\n" "HQ4x: Smooth but keeps details (best for game icons)\n" - "LANCZOS: Very smooth (best for photos)" + "LANCZOS: Very smooth (best for photos)\n" + "AI: Neural network upscaling (best quality, requires model)" ) canvas_layout.addWidget(self.upscale_method_combo)