| import glob | |
| import io | |
| import os | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import wget | |
| from torchvision.transforms import Compose, ToTensor | |
| from model import decoder, encoder | |
| WEIGHT_PATH = './weights/best_weight.pth' | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| class Model(object): | |
| def __init__(self) -> None: | |
| self.model_Enc = encoder.Encoder_RRDB(num_feat=64).to(device=DEVICE) | |
| self.model_Dec_SR = decoder.Decoder_SR_RRDB(num_in_ch=64).to(device=DEVICE) | |
| self.preprocess = Compose([ToTensor()]) | |
| self.load_model() | |
| def load_model(self, weight_path=WEIGHT_PATH): | |
| if not os.path.isfile("./weights/best_weight.pth"): | |
| response = wget.download("https://raw.githubusercontent.com/hungnguyen2611/super-resolution/master/weights/best_weight.pth", "./weights/best_weight.pth") | |
| weight = torch.load(weight_path, map_location=torch.device(DEVICE)) | |
| print("[LOADING] Loading encoder...") | |
| self.model_Enc.load_state_dict(weight['model_Enc']) | |
| print("[LOADING] Loading decoder...") | |
| self.model_Dec_SR.load_state_dict(weight['model_Dec_SR']) | |
| print("[LOADING] Loading done!") | |
| self.model_Enc.eval() | |
| self.model_Dec_SR.eval() | |
| def predict(self, img): | |
| with torch.no_grad(): | |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| img = self.preprocess(img) | |
| img = img.unsqueeze(0) | |
| img = img.to(DEVICE) | |
| feat = self.model_Enc(img) | |
| out = self.model_Dec_SR(feat) | |
| min_max = (0, 1) | |
| out = out.detach()[0].float().cpu() | |
| out = out.squeeze().float().cpu().clamp_(*min_max) | |
| out = (out - min_max[0]) / (min_max[1] - min_max[0]) | |
| out = out.numpy() | |
| out = np.transpose(out[[2, 1, 0], :, :], (1, 2, 0)) | |
| out = (out*255.0).round() | |
| out = out.astype(np.uint8) | |
| return out | |
| model = Model() | |
| def predict(img): | |
| global model | |
| img.save("test/1.png", "PNG") | |
| image = cv2.imread("test/1.png", cv2.IMREAD_COLOR) | |
| out = model.predict(img=image) | |
| cv2.imwrite(f'images_uploaded/1.png', out) | |
| return f"images_uploaded/1.png" | |
| if __name__ == '__main__': | |
| title = "Super-Resolution Demo USR-DA Unofficial 🚀🚀🔥" | |
| description = ''' | |
| <br> | |
| **This Demo expects low-quality and low-resolution images, better visual on real-world images** | |
| </br> | |
| ''' | |
| article = "<p style='text-align: center'><a href='https://openaccess.thecvf.com/content/ICCV2021/papers/Wang_Unsupervised_Real-World_Super-Resolution_A_Domain_Adaptation_Perspective_ICCV_2021_paper.pdf' target='_blank'>Unsupervised Real-World Super-Resolution: A Domain Adaptation Perspective</a> | <a href='https://github.com/hungnguyen2611/super-resolution.git' target='_blank'>Github Repo</a></p>" | |
| examples= glob.glob("testsets/*.png") | |
| gr.Interface( | |
| predict, | |
| gr.inputs.Image(type="pil", label="Input").style(height=260), | |
| gr.inputs.Image(type="pil", label="Ouput").style(height=240), | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| ).launch(enable_queue=True) | |