Update app.py

app.py

@@ -8,271 +8,14 @@ import glob
 import numpy as np
 import os
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import gradio as gr
 
+from utils.download_url import load_file_from_url
+from utils.color_fix import wavelet_reconstruction
+from models.safmn_arch import SAFMN
 
-########################################## Wavelet colorfix ###################################
-from PIL import Image
-from torch import Tensor
-from torchvision.transforms import ToTensor, ToPILImage
-def adain_color_fix(target: Image, source: Image):
-    # Convert images to tensors
-    to_tensor = ToTensor()
-    target_tensor = to_tensor(target).unsqueeze(0)
-    source_tensor = to_tensor(source).unsqueeze(0)
-
-    # Apply adaptive instance normalization
-    result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
-
-    # Convert tensor back to image
-    to_image = ToPILImage()
-    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
-
-    return result_image
-
-def wavelet_color_fix(target: Image, source: Image):
-    if target.size() != source.size():
-        source = source.resize((target.size()[-2], target.size()[-1]), Image.LANCZOS)
-    # Convert images to tensors
-    to_tensor = ToTensor()
-    target_tensor = to_tensor(target).unsqueeze(0)
-    source_tensor = to_tensor(source).unsqueeze(0)
-
-    # Apply wavelet reconstruction
-    result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
-
-    # Convert tensor back to image
-    to_image = ToPILImage()
-    result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
-
-    return result_image
-
-def calc_mean_std(feat: Tensor, eps=1e-5):
-    """Calculate mean and std for adaptive_instance_normalization.
-    Args:
-        feat (Tensor): 4D tensor.
-        eps (float): A small value added to the variance to avoid
-            divide-by-zero. Default: 1e-5.
-    """
-    size = feat.size()
-    assert len(size) == 4, 'The input feature should be 4D tensor.'
-    b, c = size[:2]
-    feat_var = feat.view(b, c, -1).var(dim=2) + eps
-    feat_std = feat_var.sqrt().view(b, c, 1, 1)
-    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
-    return feat_mean, feat_std
-
-def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
-    """Adaptive instance normalization.
-    Adjust the reference features to have the similar color and illuminations
-    as those in the degradate features.
-    Args:
-        content_feat (Tensor): The reference feature.
-        style_feat (Tensor): The degradate features.
-    """
-    size = content_feat.size()
-    style_mean, style_std = calc_mean_std(style_feat)
-    content_mean, content_std = calc_mean_std(content_feat)
-    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
-    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
-
-def wavelet_blur(image: Tensor, radius: int):
-    """
-    Apply wavelet blur to the input tensor.
-    """
-    # input shape: (1, 3, H, W)
-    # convolution kernel
-    kernel_vals = [
-        [0.0625, 0.125, 0.0625],
-        [0.125, 0.25, 0.125],
-        [0.0625, 0.125, 0.0625],
-    ]
-    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
-    # add channel dimensions to the kernel to make it a 4D tensor
-    kernel = kernel[None, None]
-    # repeat the kernel across all input channels
-    kernel = kernel.repeat(3, 1, 1, 1)
-    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
-    # apply convolution
-    output = F.conv2d(image, kernel, groups=3, dilation=radius)
-    return output
-
-def wavelet_decomposition(image: Tensor, levels=5):
-    """
-    Apply wavelet decomposition to the input tensor.
-    This function only returns the low frequency & the high frequency.
-    """
-    high_freq = torch.zeros_like(image)
-    for i in range(levels):
-        radius = 2 ** i
-        low_freq = wavelet_blur(image, radius)
-        high_freq += (image - low_freq)
-        image = low_freq
-
-    return high_freq, low_freq
-
-def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
-    """
-    Apply wavelet decomposition, so that the content will have the same color as the style.
-    """
-    # calculate the wavelet decomposition of the content feature
-    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
-    del content_low_freq
-    # calculate the wavelet decomposition of the style feature
-    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
-    del style_high_freq
-    # reconstruct the content feature with the style's high frequency
-    return content_high_freq + style_low_freq
-
-    
-########################################## URL Load  ###################################
-from torch.hub import download_url_to_file, get_dir
-from urllib.parse import urlparse
-
-def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
-    """Load file form http url, will download models if necessary.
-
-    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
-
-    Args:
-        url (str): URL to be downloaded.
-        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
-            Default: None.
-        progress (bool): Whether to show the download progress. Default: True.
-        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
-
-    Returns:
-        str: The path to the downloaded file.
-    """
-    if model_dir is None:  # use the pytorch hub_dir
-        hub_dir = get_dir()
-        model_dir = os.path.join(hub_dir, 'checkpoints')
-
-    os.makedirs(model_dir, exist_ok=True)
-
-    parts = urlparse(url)
-    filename = os.path.basename(parts.path)
-    if file_name is not None:
-        filename = file_name
-    cached_file = os.path.abspath(os.path.join(model_dir, filename))
-    if not os.path.exists(cached_file):
-        print(f'Downloading: "{url}" to {cached_file}\n')
-        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
-    return cached_file
-
-
-########################################## Model Define  ###################################
-# Layer Norm
-class LayerNorm(nn.Module):
-    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_first"):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(normalized_shape))
-        self.bias = nn.Parameter(torch.zeros(normalized_shape))
-        self.eps = eps
-        self.data_format = data_format
-        if self.data_format not in ["channels_last", "channels_first"]:
-            raise NotImplementedError
-        self.normalized_shape = (normalized_shape, )
-
-    def forward(self, x):
-        if self.data_format == "channels_last":
-            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
-        elif self.data_format == "channels_first":
-            u = x.mean(1, keepdim=True)
-            s = (x - u).pow(2).mean(1, keepdim=True)
-            x = (x - u) / torch.sqrt(s + self.eps)
-            x = self.weight[:, None, None] * x + self.bias[:, None, None]
-            return x
-
-# CCM
-class CCM(nn.Module):
-    def __init__(self, dim, growth_rate=2.0):
-        super().__init__()
-        hidden_dim = int(dim * growth_rate)
-
-        self.ccm = nn.Sequential(
-            nn.Conv2d(dim, hidden_dim, 3, 1, 1),
-            nn.GELU(), 
-            nn.Conv2d(hidden_dim, dim, 1, 1, 0)
-        )
-
-    def forward(self, x):
-        return self.ccm(x)
-
-
-# SAFM
-class SAFM(nn.Module):
-    def __init__(self, dim, n_levels=4):
-        super().__init__()
-        self.n_levels = n_levels
-        chunk_dim = dim // n_levels
-
-        # Spatial Weighting
-        self.mfr = nn.ModuleList([nn.Conv2d(chunk_dim, chunk_dim, 3, 1, 1, groups=chunk_dim) for i in range(self.n_levels)])
-        
-        # # Feature Aggregation
-        self.aggr = nn.Conv2d(dim, dim, 1, 1, 0)
-        
-        # Activation
-        self.act = nn.GELU() 
-
-    def forward(self, x):
-        h, w = x.size()[-2:]
-
-        xc = x.chunk(self.n_levels, dim=1)
-        out = []
-        for i in range(self.n_levels):
-            if i > 0:
-                p_size = (h//2**i, w//2**i)
-                s = F.adaptive_max_pool2d(xc[i], p_size)
-                s = self.mfr[i](s)
-                s = F.interpolate(s, size=(h, w), mode='nearest')
-            else:
-                s = self.mfr[i](xc[i])
-            out.append(s)
-
-        out = self.aggr(torch.cat(out, dim=1))
-        out = self.act(out) * x
-        return out
-
-class AttBlock(nn.Module):
-    def __init__(self, dim, ffn_scale=2.0):
-        super().__init__()
-
-        self.norm1 = LayerNorm(dim) 
-        self.norm2 = LayerNorm(dim) 
-
-        # Multiscale Block
-        self.safm = SAFM(dim) 
-        # Feedforward layer
-        self.ccm = CCM(dim, ffn_scale) 
-
-    def forward(self, x):
-        x = self.safm(self.norm1(x)) + x
-        x = self.ccm(self.norm2(x)) + x
-        return x
-        
-        
-class SAFMN(nn.Module):
-    def __init__(self, dim, n_blocks=8, ffn_scale=2.0, upscaling_factor=4):
-        super().__init__()
-        self.to_feat = nn.Conv2d(3, dim, 3, 1, 1)
-
-        self.feats = nn.Sequential(*[AttBlock(dim, ffn_scale) for _ in range(n_blocks)])
-
-        self.to_img = nn.Sequential(
-            nn.Conv2d(dim, 3 * upscaling_factor**2, 3, 1, 1),
-            nn.PixelShuffle(upscaling_factor)
-        )
-
-    def forward(self, x):
-        x = self.to_feat(x)
-        x = self.feats(x) + x
-        x = self.to_img(x)
-        return x
-        
+ 
 ########################################## Gradio inference  ###################################
 pretrain_model_url = {
 	'safmn_x2': 'https://github.com/sunny2109/SAFMN/releases/download/v0.1.0/SAFMN_L_Real_LSDIR_x2-v2.pth',
@@ -425,44 +168,75 @@ def inference(image, upscale, large_input_flag, color_fix):
 
 
 
-title = "Spatially-Adaptive Feature Modulation for Efficient Image Super-Resolution"
-description = r"""
-<b>Official Gradio demo</b> for <a href='https://github.com/sunny2109/SAFMN' target='_blank'><b>Spatially-Adaptive Feature Modulation for Efficient Image Super-Resolution (ICCV 2023)</b></a>.<br>
-"""
-article = r"""
-If SAFMN is helpful, please help to ⭐ the <a href='https://github.com/sunny2109/SAFMN' target='_blank'>Github Repo</a>. Thanks!
-[![GitHub Stars](https://img.shields.io/github/stars/sunny2109/SAFMN?style=social)](https://github.com/sunny2109/SAFMN)
+title = "SAFMN for Real-world SR"
+description = ''' ### Spatially-Adaptive Feature Modulation for Efficient Image Super-Resolution - ICCV 2023
+#### Long Sun](https://github.com/sunny2109), [Jiangxin Dong](https://scholar.google.com/citations?user=ruebFVEAAAAJ&hl=zh-CN&oi=ao), [Jinhui Tang](https://scholar.google.com/citations?user=ByBLlEwAAAAJ&hl=zh-CN), and [Jinshan Pan](https://jspan.github.io/)
+#### [IMAG Lab](https://imag-njust.net/), Nanjing University of Science and Technology
+
 
----
-📝 **Citation**
+####  Drag the slider on the super-resolution image left and right to see the changes in the image details. SAFMN performs x2/x4 upscaling on the input image.
+<br>
 
-If our work is useful for your research, please consider citing:
-```bibtex
+### If our work is useful for your research, please consider citing:
+<code>
 @inproceedings{sun2023safmn,
     title={Spatially-Adaptive Feature Modulation for Efficient Image Super-Resolution},
     author={Sun, Long and Dong, Jiangxin and Tang, Jinhui and Pan, Jinshan},
-    booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
+    booktitle={ICCV},
     year={2023}
 }
-```
-
-<center><img src='https://visitor-badge.laobi.icu/badge?page_id=sunny2109/SAFMN' alt='visitors'></center>
+</code>
+<br>
+'''
+
+
+article = "<p style='text-align: center'><a href='https://eduardzamfir.github.io/seemore' target='_blank'>See More Details: Efficient Image Super-Resolution by Experts Mining</a></p>"
+
+#### Image,Prompts examples
+examples = [
+            ['images/0801x4.png'],
+            ['images/0840x4.png'],
+            ['images/0841x4.png'],
+            ['images/0870x4.png'],
+            ['images/0878x4.png'],
+            ['images/0884x4.png'],
+            ['images/0900x4.png'],
+            ['images/img002x4.png'],
+            ['images/img003x4.png'],
+            ['images/img004x4.png'],
+            ['images/img035x4.png'],
+            ['images/img053x4.png'],
+            ['images/img064x4.png'],
+            ['images/img083x4.png'],
+            ['images/img092x4.png'],
+            ]
+
+css = """
+    .image-frame img, .image-container img {
+        width: auto;
+        height: auto;
+        max-width: none;
+    }
 """
 
 demo = gr.Interface(
-    inference, [
-        gr.inputs.Image(type="filepath", label="Input"),
-        gr.inputs.Number(default=2, label="Upscaling factor (up to 4)"),
-		gr.inputs.Checkbox(default=False, label="Memory-efficient inference"),
-        gr.inputs.Checkbox(default=False, label="Color correction"),
-    ], [
-        gr.outputs.Image(type="numpy", label="Output"),
-        gr.outputs.File(label="Download the output")
+    fn=process_img,
+    inputs=[
+        gr.Image(type="pil", label="Input", value="real_testdata/004.png"),
+        gr.Number(default=2, label="Upscaling factor (up to 4)"),
+		gr.Checkbox(default=False, label="Memory-efficient inference"),
+        gr.Checkbox(default=False, label="Color correction"),
     ],
+    outputs=ImageSlider(label="Super-Resolved Image", 
+                        type="pil",
+                        show_download_button=True,
+                        ),
     title=title,
     description=description,
-    article=article,       
+    article=article,
+    examples=examples,
+    css=css,
 )
 
-demo.queue(concurrency_count=2)
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()