Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchvision.io import read_image | |
| from torchvision.utils import save_image | |
| import os | |
| from torchvision.transforms.functional import to_tensor, to_pil_image | |
| # === RCAN MODULES === | |
| class CALayer(nn.Module): | |
| def __init__(self, channels, reduction=16): | |
| super().__init__() | |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) | |
| self.fc = nn.Sequential( | |
| nn.Conv2d(channels, channels // reduction, 1), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(channels // reduction, channels, 1), | |
| nn.Sigmoid() | |
| ) | |
| def forward(self, x): | |
| w = self.fc(self.avg_pool(x)) | |
| return x * w | |
| class RCAB(nn.Module): | |
| def __init__(self, channels): | |
| super().__init__() | |
| self.body = nn.Sequential( | |
| nn.Conv2d(channels, channels, 3, padding=1), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(channels, channels, 3, padding=1), | |
| CALayer(channels) | |
| ) | |
| def forward(self, x): | |
| return x + self.body(x) | |
| class ResidualGroup(nn.Module): | |
| def __init__(self, channels, n_rcab): | |
| super().__init__() | |
| modules = [RCAB(channels) for _ in range(n_rcab)] | |
| modules.append(nn.Conv2d(channels, channels, 3, padding=1)) | |
| self.body = nn.Sequential(*modules) | |
| def forward(self, x): | |
| return x + self.body(x) | |
| class Upsampler(nn.Sequential): | |
| def __init__(self, scale, channels): | |
| m = [] | |
| for _ in range(int(torch.log2(torch.tensor(scale)))): | |
| m.append(nn.Conv2d(channels, channels * 4, 3, padding=1)) | |
| m.append(nn.PixelShuffle(2)) | |
| super().__init__(*m) | |
| class RCAN(nn.Module): | |
| def __init__(self, in_channels=3, out_channels=3, n_feat=64, n_rg=10, n_rcab=20, scale=4): | |
| super().__init__() | |
| self.head = nn.Conv2d(in_channels, n_feat, 3, padding=1) | |
| self.body = nn.Sequential( | |
| *[ResidualGroup(n_feat, n_rcab) for _ in range(n_rg)], | |
| nn.Conv2d(n_feat, n_feat, 3, padding=1) | |
| ) | |
| self.upsample = Upsampler(scale, n_feat) | |
| self.tail = nn.Conv2d(n_feat, out_channels, 3, padding=1) | |
| def forward(self, x): | |
| x = self.head(x) | |
| res = self.body(x) | |
| x = x + res | |
| x = self.upsample(x) | |
| return self.tail(x) | |
| # === INFERENCE === | |
| def rcan_upscale(lr_img_pil, model_path="weights/rcan_epoch_20.pth", device='cpu'): | |
| """ | |
| Super resolves a low-resolution PIL image using the RCAN model. | |
| Args: | |
| lr_img_pil (PIL.Image): Low-resolution input image. | |
| model_path (str): Path to the model weights. | |
| device (str): 'cuda' or 'cpu'. | |
| Returns: | |
| PIL.Image: High-resolution output image. | |
| """ | |
| # Load model | |
| device = torch.device(device if torch.cuda.is_available() else 'cpu') | |
| model = RCAN(scale=4) | |
| model.load_state_dict(torch.load(model_path, map_location=device)) | |
| model.to(device).eval() | |
| # Convert PIL image to normalized tensor | |
| lr_tensor = to_tensor(lr_img_pil).unsqueeze(0).to(device) # Add batch dim | |
| # Inference | |
| with torch.no_grad(): | |
| sr_tensor = model(lr_tensor).squeeze(0).clamp(0, 1).cpu() # Remove batch | |
| # Convert tensor back to PIL image | |
| sr_img_pil = to_pil_image(sr_tensor) | |
| return sr_img_pil | |