Eric Rollei
d6dd6470cf
Features: - 40+ image processing algorithms - Traditional denoising: Wavelet, NLM, BM3D (GPU-accelerated) - AI models: DnCNN, SwinIR, SCUNet, NAFNet, Noise-DA - DeepInv diffusion denoising integration - Advanced film grain detection and processing - Frequency domain processing - Richardson-Lucy deconvolution - Wiener filter - External DeepInv service with FastAPI - Comprehensive documentation and guides
428 lines
16 KiB
Python
428 lines
16 KiB
Python
"""
|
|
Real-ESRGAN Processing Implementation
|
|
Based on 2024-2025 research findings for practical super-resolution
|
|
|
|
Real-ESRGAN variants:
|
|
- RealESRGAN_x4plus: General scenes
|
|
- RealESRGAN_x4plus_anime_6B: Illustrations/anime
|
|
- realesr-general-x4v3: Lightweight alternative
|
|
|
|
Reference: https://github.com/xinntao/Real-ESRGAN
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import cv2
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional, Tuple
|
|
import warnings
|
|
import logging
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ResidualDenseBlock(nn.Module):
|
|
"""Residual Dense Block for ESRGAN"""
|
|
def __init__(self, num_feat=64, num_grow_ch=32):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
|
|
self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
|
|
def forward(self, x):
|
|
x1 = self.lrelu(self.conv1(x))
|
|
x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
|
|
x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
|
|
x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
|
|
x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
|
|
return x5 * 0.2 + x
|
|
|
|
class RRDB(nn.Module):
|
|
"""Residual in Residual Dense Block"""
|
|
def __init__(self, num_feat, num_grow_ch=32):
|
|
super().__init__()
|
|
self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
|
|
self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
|
|
self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
|
|
|
|
def forward(self, x):
|
|
out = self.rdb1(x)
|
|
out = self.rdb2(out)
|
|
out = self.rdb3(out)
|
|
return out * 0.2 + x
|
|
|
|
class RealESRGANProcessor:
|
|
"""Real-ESRGAN processor for practical super-resolution and enhancement"""
|
|
|
|
def __init__(self, model_name: str = "RealESRGAN_x4plus", device: str = "auto"):
|
|
"""
|
|
Initialize Real-ESRGAN processor
|
|
|
|
Args:
|
|
model_name: Model variant to use
|
|
device: Processing device (auto, cpu, cuda)
|
|
"""
|
|
self.model_name = model_name
|
|
self.device = self._get_device(device)
|
|
self.model = None
|
|
self.scale = 4 # Default scale factor
|
|
|
|
# Check if model file exists
|
|
models_dir = Path(__file__).parent.parent / "models"
|
|
model_path = models_dir / f"{model_name}.pth"
|
|
if model_path.exists():
|
|
logger.info(f"Found pretrained model: {model_name}.pth")
|
|
self.model_path = str(model_path)
|
|
else:
|
|
logger.warning(f"Model {model_name}.pth not found in models directory")
|
|
logger.info("Download models using: python download_models.py")
|
|
self.model_path = None
|
|
|
|
# Model configurations
|
|
self.model_configs = {
|
|
"RealESRGAN_x4plus": {
|
|
"scale": 4,
|
|
"num_in_ch": 3,
|
|
"num_out_ch": 3,
|
|
"num_feat": 64,
|
|
"num_block": 23,
|
|
"num_grow_ch": 32,
|
|
"description": "General scenes, best overall quality"
|
|
},
|
|
"RealESRGAN_x4plus_anime_6B": {
|
|
"scale": 4,
|
|
"num_in_ch": 3,
|
|
"num_out_ch": 3,
|
|
"num_feat": 64,
|
|
"num_block": 6,
|
|
"num_grow_ch": 32,
|
|
"description": "Optimized for illustrations and anime"
|
|
},
|
|
"realesr-general-x4v3": {
|
|
"scale": 4,
|
|
"num_in_ch": 3,
|
|
"num_out_ch": 3,
|
|
"num_feat": 64,
|
|
"num_block": 16,
|
|
"num_grow_ch": 32,
|
|
"description": "Lightweight alternative"
|
|
}
|
|
}
|
|
|
|
# Tile processing for memory efficiency
|
|
self.tile_size = 512
|
|
self.tile_pad = 32
|
|
|
|
def _get_device(self, device: str) -> torch.device:
|
|
"""Get appropriate device for processing"""
|
|
if device == "auto":
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
else:
|
|
return torch.device("cpu")
|
|
return torch.device(device)
|
|
|
|
def _load_model(self) -> bool:
|
|
"""Load or create Real-ESRGAN model"""
|
|
try:
|
|
# Try to import Real-ESRGAN
|
|
try:
|
|
from realesrgan import RealESRGANer
|
|
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
|
|
|
|
# Use official Real-ESRGAN implementation
|
|
model_path = self.model_path # Use downloaded model if available
|
|
self.model = RealESRGANer(
|
|
scale=self.scale,
|
|
model_path=model_path,
|
|
model=None,
|
|
tile=self.tile_size,
|
|
tile_pad=self.tile_pad,
|
|
pre_pad=0,
|
|
half=True if self.device.type == 'cuda' else False,
|
|
gpu_id=0 if self.device.type == 'cuda' else None
|
|
)
|
|
logger.info(f"Loaded Real-ESRGAN model: {self.model_name}")
|
|
return True
|
|
|
|
except ImportError:
|
|
logger.warning("Real-ESRGAN not installed, using fallback implementation")
|
|
return self._create_fallback_model()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load Real-ESRGAN model: {e}")
|
|
return self._create_fallback_model()
|
|
|
|
def _create_fallback_model(self) -> bool:
|
|
"""Create fallback super-resolution model"""
|
|
try:
|
|
# Simple ESRGAN-style architecture as fallback
|
|
config = self.model_configs[self.model_name]
|
|
|
|
class FallbackESRGAN(nn.Module):
|
|
def __init__(self, num_in_ch=3, num_out_ch=3, scale=4, num_feat=64, num_block=16, num_grow_ch=32):
|
|
super().__init__()
|
|
self.scale = scale
|
|
|
|
# Feature extraction
|
|
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
|
|
|
|
# Residual blocks
|
|
self.trunk_conv = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
|
|
self.body = nn.Sequential(*[
|
|
ResidualDenseBlock(num_feat, num_grow_ch) for _ in range(num_block)
|
|
])
|
|
|
|
# Upsampling
|
|
self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
|
|
self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
|
|
self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
|
|
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
|
|
def forward(self, x):
|
|
feat = self.lrelu(self.conv_first(x))
|
|
|
|
# Trunk
|
|
trunk = self.trunk_conv(self.body(feat))
|
|
feat = feat + trunk
|
|
|
|
# Upsampling
|
|
feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
|
|
feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
|
|
out = self.conv_last(self.lrelu(self.conv_hr(feat)))
|
|
|
|
return out
|
|
|
|
# Create model
|
|
model_config = {k: v for k, v in config.items() if k != 'description'}
|
|
self.model = FallbackESRGAN(**model_config).to(self.device)
|
|
|
|
# Load pretrained weights if available
|
|
if self.model_path and Path(self.model_path).exists():
|
|
try:
|
|
checkpoint = torch.load(self.model_path, map_location='cpu')
|
|
|
|
# Handle different checkpoint formats
|
|
if 'params' in checkpoint:
|
|
state_dict = checkpoint['params']
|
|
elif 'params_ema' in checkpoint:
|
|
state_dict = checkpoint['params_ema']
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
# Load state dict
|
|
self.model.load_state_dict(state_dict, strict=False)
|
|
logger.info(f"Loaded pretrained weights from {self.model_path}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load pretrained weights: {e}")
|
|
logger.info("Using random initialization")
|
|
else:
|
|
logger.info("No pretrained weights found, using random initialization")
|
|
|
|
self.model.eval()
|
|
|
|
logger.info(f"Created fallback Real-ESRGAN model: {self.model_name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to create fallback model: {e}")
|
|
return False
|
|
|
|
def process_image(self, image: np.ndarray, **kwargs) -> np.ndarray:
|
|
"""
|
|
Process image with Real-ESRGAN
|
|
|
|
Args:
|
|
image: Input image (H, W, C) in RGB format
|
|
**kwargs: Additional processing parameters
|
|
|
|
Returns:
|
|
Enhanced image
|
|
"""
|
|
if self.model is None:
|
|
if not self._load_model():
|
|
return self._fallback_bicubic(image)
|
|
|
|
try:
|
|
# Prepare image
|
|
if image.dtype != np.uint8:
|
|
image = (image * 255).astype(np.uint8)
|
|
|
|
logger.info(f"Input image shape: {image.shape}, dtype: {image.dtype}, range: [{image.min()}, {image.max()}]")
|
|
|
|
# Process with Real-ESRGAN
|
|
if hasattr(self.model, 'enhance'):
|
|
# Official Real-ESRGAN
|
|
logger.info("Using official Real-ESRGAN enhance method")
|
|
output, _ = self.model.enhance(image, outscale=self.scale)
|
|
logger.info(f"Raw output shape: {output.shape}, dtype: {output.dtype}, range: [{output.min()}, {output.max()}]")
|
|
|
|
# Ensure output is in correct format
|
|
if output.dtype == np.uint8:
|
|
# Convert uint8 [0,255] to float32 [0,1]
|
|
output = output.astype(np.float32) / 255.0
|
|
elif output.dtype in [np.float32, np.float64]:
|
|
# Check if it's already in [0,1] range
|
|
if output.max() <= 1.0:
|
|
output = output.astype(np.float32)
|
|
else:
|
|
# Assume it's in [0,255] range
|
|
output = output.astype(np.float32) / 255.0
|
|
|
|
logger.info(f"Final output shape: {output.shape}, dtype: {output.dtype}, range: [{output.min():.3f}, {output.max():.3f}]")
|
|
|
|
else:
|
|
# Fallback model
|
|
logger.info("Using fallback model")
|
|
output = self._process_with_fallback(image)
|
|
|
|
return output
|
|
|
|
except Exception as e:
|
|
logger.error(f"Real-ESRGAN processing failed: {e}")
|
|
return self._fallback_bicubic(image)
|
|
|
|
def _process_with_fallback(self, image: np.ndarray) -> np.ndarray:
|
|
"""Process image with fallback model"""
|
|
# Convert to tensor
|
|
img_tensor = torch.from_numpy(image).float().permute(2, 0, 1).unsqueeze(0) / 255.0
|
|
img_tensor = img_tensor.to(self.device)
|
|
|
|
# Process
|
|
with torch.no_grad():
|
|
output = self.model(img_tensor)
|
|
|
|
# Convert back
|
|
output = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
|
output = np.clip(output, 0, 1)
|
|
|
|
# Normalize to use full range if output is too dark
|
|
if output.max() < 0.8: # If max is less than 0.8, normalize to full range
|
|
output_min = output.min()
|
|
output_max = output.max()
|
|
if output_max > output_min:
|
|
output = (output - output_min) / (output_max - output_min)
|
|
logger.info(f"Normalized fallback output from [{output_min:.3f}, {output_max:.3f}] to [0, 1]")
|
|
|
|
return (output * 255).astype(np.uint8)
|
|
|
|
def _fallback_bicubic(self, image: np.ndarray) -> np.ndarray:
|
|
"""Fallback bicubic upsampling"""
|
|
h, w = image.shape[:2]
|
|
new_h, new_w = h * self.scale, w * self.scale
|
|
|
|
upscaled = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
|
|
|
|
# Apply some basic enhancement
|
|
# Unsharp masking
|
|
gaussian = cv2.GaussianBlur(upscaled, (0, 0), 2.0)
|
|
upscaled = cv2.addWeighted(upscaled, 1.5, gaussian, -0.5, 0)
|
|
|
|
return upscaled
|
|
|
|
def get_model_info(self) -> Dict[str, Any]:
|
|
"""Get model information"""
|
|
config = self.model_configs.get(self.model_name, {})
|
|
return {
|
|
"model_name": self.model_name,
|
|
"scale_factor": self.scale,
|
|
"device": str(self.device),
|
|
"tile_size": self.tile_size,
|
|
"description": config.get("description", "Unknown model"),
|
|
"memory_usage": "4-12GB VRAM (depending on input size)",
|
|
"supported_formats": ["RGB", "BGR"],
|
|
"input_range": "[0, 1] float or [0, 255] uint8"
|
|
}
|
|
|
|
|
|
class ResidualDenseBlock(nn.Module):
|
|
"""Residual Dense Block for ESRGAN"""
|
|
|
|
def __init__(self, num_feat=64, num_grow_ch=32):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
|
|
self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
|
|
self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
|
|
def forward(self, x):
|
|
x1 = self.lrelu(self.conv1(x))
|
|
x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
|
|
x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
|
|
x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
|
|
x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
|
|
|
|
return x5 * 0.2 + x
|
|
|
|
|
|
def get_realesrgan_presets() -> Dict[str, Dict[str, Any]]:
|
|
"""Get available Real-ESRGAN presets"""
|
|
return {
|
|
"general": {
|
|
"model_name": "RealESRGAN_x4plus",
|
|
"description": "Best for general photographs and natural scenes",
|
|
"scale": 4,
|
|
"memory_usage": "8-12GB VRAM"
|
|
},
|
|
"anime": {
|
|
"model_name": "RealESRGAN_x4plus_anime_6B",
|
|
"description": "Optimized for illustrations, anime, and artistic images",
|
|
"scale": 4,
|
|
"memory_usage": "4-8GB VRAM"
|
|
},
|
|
"lightweight": {
|
|
"model_name": "realesr-general-x4v3",
|
|
"description": "Lightweight version for faster processing",
|
|
"scale": 4,
|
|
"memory_usage": "4-6GB VRAM"
|
|
}
|
|
}
|
|
|
|
|
|
def process_with_realesrgan(image: np.ndarray,
|
|
model_name: str = "RealESRGAN_x4plus",
|
|
device: str = "auto",
|
|
**kwargs) -> Tuple[np.ndarray, Dict[str, Any]]:
|
|
"""
|
|
Process image with Real-ESRGAN
|
|
|
|
Args:
|
|
image: Input image (H, W, C) in RGB format
|
|
model_name: Model variant to use
|
|
device: Processing device
|
|
**kwargs: Additional parameters
|
|
|
|
Returns:
|
|
(enhanced_image, processing_info)
|
|
"""
|
|
processor = RealESRGANProcessor(model_name, device)
|
|
enhanced = processor.process_image(image, **kwargs)
|
|
info = processor.get_model_info()
|
|
|
|
return enhanced, info
|
|
|
|
|
|
# Example usage and testing
|
|
if __name__ == "__main__":
|
|
# Test with synthetic image
|
|
test_image = np.random.rand(256, 256, 3).astype(np.float32)
|
|
|
|
print("Testing Real-ESRGAN processing...")
|
|
enhanced, info = process_with_realesrgan(test_image)
|
|
|
|
print(f"Original shape: {test_image.shape}")
|
|
print(f"Enhanced shape: {enhanced.shape}")
|
|
print(f"Model info: {info}")
|