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
431 lines
16 KiB
Python
431 lines
16 KiB
Python
"""
|
||
Restormer transformer-based restoration processor with tiling and sharpness controls.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import os
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Dict, Iterable, List, Optional, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
try:
|
||
from ..network_restormer import Restormer
|
||
except ImportError: # pragma: no cover - allow standalone execution
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
sys.path.append(str(Path(__file__).parent.parent))
|
||
from network_restormer import Restormer
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ModelConfig:
|
||
identifier: str
|
||
scale: int
|
||
filename_candidates: Tuple[str, ...]
|
||
download_url: Optional[str]
|
||
description: str
|
||
|
||
|
||
MODEL_CONFIGS: Dict[str, ModelConfig] = {
|
||
"real_sr_x4": ModelConfig(
|
||
identifier="real_sr_x4",
|
||
scale=4,
|
||
filename_candidates=(
|
||
"Restormer_RealSR_DFOWM_x4.pth",
|
||
"Restormer_RealSR_DFOWM_x4_GAN.pth",
|
||
"Restormer_RealSR_DFOWM_x4_PSNR.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_RealSR_DFOWM_x4.pth",
|
||
description="Real-world super-resolution x4 (DFOWM dataset)",
|
||
),
|
||
"denoise_sigma15": ModelConfig(
|
||
identifier="denoise_sigma15",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Denoising_Color_sigma15.pth",
|
||
"Restormer_GaussianColorSigma15.pth",
|
||
"gaussian_color_denoising_sigma15.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Denoising_Color_sigma15.pth",
|
||
description="Gaussian color denoising σ=15",
|
||
),
|
||
"denoise_sigma25": ModelConfig(
|
||
identifier="denoise_sigma25",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Denoising_Color_sigma25.pth",
|
||
"Restormer_GaussianColorSigma25.pth",
|
||
"gaussian_color_denoising_sigma25.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Denoising_Color_sigma25.pth",
|
||
description="Gaussian color denoising σ=25",
|
||
),
|
||
"denoise_sigma50": ModelConfig(
|
||
identifier="denoise_sigma50",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Denoising_Color_sigma50.pth",
|
||
"Restormer_GaussianColorSigma50.pth",
|
||
"gaussian_color_denoising_sigma50.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Denoising_Color_sigma50.pth",
|
||
description="Gaussian color denoising σ=50",
|
||
),
|
||
"real_denoise": ModelConfig(
|
||
identifier="real_denoise",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_RealDenoising.pth",
|
||
"Restormer_SIDD.pth",
|
||
"real_denoising.pth",
|
||
"restormer_real_denoising.pth",
|
||
"gaussian_color_denoising_blind.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_RealDenoising.pth",
|
||
description="Real image denoising (SIDD)",
|
||
),
|
||
"motion_deblur": ModelConfig(
|
||
identifier="motion_deblur",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_MotionDeblurring.pth",
|
||
"Restormer_GoPro.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_MotionDeblurring.pth",
|
||
description="Motion deblurring (GoPro)",
|
||
),
|
||
"defocus_deblur": ModelConfig(
|
||
identifier="defocus_deblur",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_DefocusDeblurring_x4.pth",
|
||
"Restormer_DPDD.pth",
|
||
"dual_pixel_defocus_deblurring.pth",
|
||
"single_image_defocus_deblurring.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_DefocusDeblurring_x4.pth",
|
||
description="Defocus deblurring (DPDD)",
|
||
),
|
||
"derain_indoor": ModelConfig(
|
||
identifier="derain_indoor",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Deraining_Indoor.pth",
|
||
"Restormer_Rain100H.pth",
|
||
"deraining.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Deraining_Indoor.pth",
|
||
description="Deraining for heavy indoor streaks",
|
||
),
|
||
"derain_outdoor": ModelConfig(
|
||
identifier="derain_outdoor",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Deraining_Outdoor.pth",
|
||
"Restormer_SPA-Data.pth",
|
||
"deraining.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Deraining_Outdoor.pth",
|
||
description="Deraining for outdoor scenes (SPA-Data)",
|
||
),
|
||
"dehaze": ModelConfig(
|
||
identifier="dehaze",
|
||
scale=1,
|
||
filename_candidates=(
|
||
"Restormer_Dehazing.pth",
|
||
"Restormer_ITS.pth",
|
||
),
|
||
download_url="https://github.com/swz30/Restormer/releases/download/v1.0/Restormer_Dehazing.pth",
|
||
description="Dehazing (ITS dataset)",
|
||
),
|
||
}
|
||
|
||
|
||
class RestormerProcessor:
|
||
"""Wrapper around Restormer with tiling support and detail boosting."""
|
||
|
||
def __init__(
|
||
self,
|
||
task: str = "real_sr_x4",
|
||
device: str = "auto",
|
||
models_dir: Optional[str] = None,
|
||
tile_size: Optional[int] = None,
|
||
tile_overlap: int = 48,
|
||
) -> None:
|
||
if task not in MODEL_CONFIGS:
|
||
raise ValueError(f"Unknown Restormer task: {task}")
|
||
|
||
self.config = MODEL_CONFIGS[task]
|
||
self.scale = self.config.scale
|
||
self.tile_size = tile_size
|
||
self.tile_overlap = tile_overlap
|
||
|
||
if device == "auto":
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
if device == "cuda" and not torch.cuda.is_available():
|
||
print("CUDA requested but not available. Falling back to CPU.")
|
||
device = "cpu"
|
||
self.device = device
|
||
|
||
default_model_dir = Path(__file__).resolve().parent.parent / "models"
|
||
self.models_dir = Path(models_dir).expanduser().resolve() if models_dir else default_model_dir
|
||
# Cache candidate directories and lazy-load model weights on this instance.
|
||
self._weight_search_dirs = self._build_weight_search_dirs(self.models_dir)
|
||
self._weights_loaded = False
|
||
self._last_weight_error: Optional[str] = None
|
||
self.model = self._load_model().to(self.device)
|
||
self.model.eval()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Model loading helpers
|
||
# ------------------------------------------------------------------
|
||
def _build_weight_search_dirs(self, base_dir: Path) -> Tuple[Path, ...]:
|
||
"""Construct candidate directories for locating Restormer weights."""
|
||
repo_root = Path(__file__).resolve().parent.parent
|
||
candidates = [base_dir]
|
||
|
||
# Support nested pretrained_weights directory alongside base dir
|
||
if base_dir.name != "pretrained_weights":
|
||
candidates.append(base_dir / "pretrained_weights")
|
||
|
||
# Fall back to project-level models directories
|
||
repo_models = repo_root / "models"
|
||
if repo_models not in candidates:
|
||
candidates.append(repo_models)
|
||
repo_models_pretrained = repo_models / "pretrained_weights"
|
||
if repo_models_pretrained not in candidates:
|
||
candidates.append(repo_models_pretrained)
|
||
|
||
# Remove duplicates while preserving order
|
||
unique_candidates = []
|
||
for path in candidates:
|
||
resolved = path.resolve()
|
||
if resolved not in unique_candidates:
|
||
unique_candidates.append(resolved)
|
||
return tuple(unique_candidates)
|
||
|
||
def _load_model(self) -> Restormer:
|
||
model = Restormer(inp_channels=3, out_channels=3)
|
||
weight_path = self._resolve_weight_path()
|
||
|
||
if weight_path is None:
|
||
message = (
|
||
f"Restormer weights for task '{self.config.identifier}' were not found.\n"
|
||
f"Expected one of: {self.config.filename_candidates}.\n"
|
||
"Place the file in one of the following directories: \n"
|
||
+ "\n".join(f" • {path}" for path in self._weight_search_dirs)
|
||
+ "\n"
|
||
f"Suggested download URL: {self.config.download_url}"
|
||
)
|
||
self._last_weight_error = message
|
||
self._weights_loaded = False
|
||
print(message)
|
||
return model
|
||
|
||
try:
|
||
checkpoint = torch.load(weight_path, map_location="cpu")
|
||
state_dict = self._extract_state_dict(checkpoint)
|
||
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
||
if missing:
|
||
print(f"Warning: missing parameters while loading Restormer weights: {missing}")
|
||
if unexpected:
|
||
print(f"Warning: unexpected parameters while loading Restormer weights: {unexpected}")
|
||
print(f"Loaded Restormer weights from {weight_path}")
|
||
self._weights_loaded = True
|
||
self._last_weight_error = None
|
||
except Exception as exc: # pragma: no cover - runtime guard
|
||
print(f"Failed to load Restormer weights: {exc}")
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
self._last_weight_error = str(exc)
|
||
self._weights_loaded = False
|
||
return model
|
||
|
||
def _resolve_weight_path(self) -> Optional[str]:
|
||
for directory in self._weight_search_dirs:
|
||
if not directory.exists():
|
||
continue
|
||
for candidate in self.config.filename_candidates:
|
||
candidate_path = directory / candidate
|
||
if candidate_path.exists():
|
||
return str(candidate_path)
|
||
return None
|
||
|
||
@staticmethod
|
||
def _extract_state_dict(checkpoint: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||
if isinstance(checkpoint, dict):
|
||
for key in ("params", "state_dict", "model", "net_g"):
|
||
if key in checkpoint and isinstance(checkpoint[key], dict):
|
||
return checkpoint[key]
|
||
return checkpoint
|
||
|
||
# ------------------------------------------------------------------
|
||
# Processing
|
||
# ------------------------------------------------------------------
|
||
def process_image(
|
||
self,
|
||
image: np.ndarray,
|
||
strength: float = 1.0,
|
||
detail_boost: float = 0.0,
|
||
detail_radius: float = 1.2,
|
||
) -> np.ndarray:
|
||
if not self._weights_loaded:
|
||
error_message = self._last_weight_error or "Restormer weights are not loaded."
|
||
raise RuntimeError(error_message)
|
||
|
||
# Support optional batch dimension
|
||
if image.ndim == 4:
|
||
if image.shape[0] != 1:
|
||
raise ValueError("Restormer processor expects a single image (batch size 1)")
|
||
image = image[0]
|
||
|
||
if image.ndim != 3 or image.shape[2] not in (3, 4):
|
||
raise ValueError("Restormer expects RGB images with 3 channels")
|
||
|
||
# Drop alpha channel if present
|
||
if image.shape[2] == 4:
|
||
image = image[..., :3]
|
||
|
||
image = image.astype(np.float32)
|
||
if image.max() > 1.0:
|
||
image /= 255.0
|
||
|
||
original = image.copy()
|
||
tensor = torch.from_numpy(image.transpose(2, 0, 1)).unsqueeze(0).to(self.device)
|
||
|
||
with torch.no_grad():
|
||
if self.tile_size is None:
|
||
output = self._forward(tensor)
|
||
else:
|
||
output = self._forward_with_tiling(tensor)
|
||
|
||
result = output.squeeze(0).cpu().numpy().transpose(1, 2, 0)
|
||
result = np.clip(result, 0.0, 1.0)
|
||
|
||
if self.scale > 1:
|
||
# Upscale original for blending when performing super-resolution
|
||
from scipy.ndimage import zoom
|
||
|
||
upscaled = zoom(original, (self.scale, self.scale, 1), order=3)
|
||
original = upscaled
|
||
|
||
result = original * (1.0 - strength) + result * strength
|
||
|
||
if detail_boost > 0:
|
||
detail_boost = float(np.clip(detail_boost, 0.0, 1.0))
|
||
sigma = max(0.1, float(detail_radius))
|
||
base = np.ascontiguousarray(result)
|
||
blurred = cv2.GaussianBlur(base, (0, 0), sigma)
|
||
blurred = np.ascontiguousarray(blurred)
|
||
detail = np.ascontiguousarray(base - blurred)
|
||
result = np.ascontiguousarray(base + detail_boost * detail)
|
||
|
||
result = np.clip(result, 0.0, 1.0)
|
||
result = np.ascontiguousarray(result)
|
||
return result
|
||
|
||
def _forward(self, tensor: torch.Tensor) -> torch.Tensor:
|
||
return self.model(tensor)
|
||
|
||
def _forward_with_tiling(self, tensor: torch.Tensor) -> torch.Tensor:
|
||
tile = int(self.tile_size)
|
||
overlap = int(self.tile_overlap)
|
||
if tile <= overlap:
|
||
raise ValueError("tile_size must be larger than tile_overlap")
|
||
|
||
tensor, pad_h, pad_w = self._pad_to_tile(tensor, tile)
|
||
b, c, h, w = tensor.shape
|
||
stride = tile - overlap
|
||
|
||
h_steps = self._compute_steps(h, tile, stride)
|
||
w_steps = self._compute_steps(w, tile, stride)
|
||
|
||
if self.scale > 1:
|
||
out_shape = (b, c, h * self.scale, w * self.scale)
|
||
else:
|
||
out_shape = (b, c, h, w)
|
||
|
||
dtype = tensor.dtype
|
||
output_accum = torch.zeros(out_shape, device=self.device, dtype=dtype)
|
||
weight_accum = torch.zeros(out_shape, device=self.device, dtype=dtype)
|
||
|
||
for y in h_steps:
|
||
for x in w_steps:
|
||
patch = tensor[:, :, y : y + tile, x : x + tile]
|
||
patch_out = self._forward(patch)
|
||
|
||
if self.scale > 1:
|
||
ys = y * self.scale
|
||
xs = x * self.scale
|
||
ye = ys + patch_out.shape[2]
|
||
xe = xs + patch_out.shape[3]
|
||
else:
|
||
ys, ye = y, y + patch_out.shape[2]
|
||
xs, xe = x, x + patch_out.shape[3]
|
||
|
||
output_accum[:, :, ys:ye, xs:xe] += patch_out
|
||
weight_accum[:, :, ys:ye, xs:xe] += 1.0
|
||
|
||
output = output_accum / torch.clamp(weight_accum, min=1e-6)
|
||
|
||
if pad_h or pad_w:
|
||
if self.scale > 1:
|
||
cut_h = (tensor.shape[2] - pad_h) * self.scale
|
||
cut_w = (tensor.shape[3] - pad_w) * self.scale
|
||
else:
|
||
cut_h = tensor.shape[2] - pad_h
|
||
cut_w = tensor.shape[3] - pad_w
|
||
output = output[:, :, :cut_h, :cut_w]
|
||
return output
|
||
|
||
@staticmethod
|
||
def _compute_steps(length: int, tile: int, stride: int) -> Iterable[int]:
|
||
positions: List[int] = []
|
||
pos = 0
|
||
while True:
|
||
positions.append(pos)
|
||
if pos + tile >= length:
|
||
break
|
||
pos += stride
|
||
if pos + tile > length:
|
||
pos = length - tile
|
||
return positions
|
||
|
||
@staticmethod
|
||
def _pad_to_tile(tensor: torch.Tensor, tile: int) -> Tuple[torch.Tensor, int, int]:
|
||
b, c, h, w = tensor.shape
|
||
pad_h = (tile - h % tile) % tile
|
||
pad_w = (tile - w % tile) % tile
|
||
if pad_h or pad_w:
|
||
tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")
|
||
return tensor, pad_h, pad_w
|
||
|
||
# ------------------------------------------------------------------
|
||
# Convenience interface
|
||
# ------------------------------------------------------------------
|
||
def __call__(
|
||
self,
|
||
image: np.ndarray,
|
||
strength: float = 1.0,
|
||
detail_boost: float = 0.0,
|
||
detail_radius: float = 1.2,
|
||
) -> np.ndarray:
|
||
return self.process_image(
|
||
image,
|
||
strength=strength,
|
||
detail_boost=detail_boost,
|
||
detail_radius=detail_radius,
|
||
)
|
||
|
||
|
||
__all__ = ["RestormerProcessor", "MODEL_CONFIGS"]
|