|
''' |
|
Author: Egrt |
|
Date: 2022-04-07 14:00:52 |
|
LastEditors: [egrt] |
|
LastEditTime: 2022-05-04 11:47:21 |
|
FilePath: \MaskGAN\maskgan.py |
|
''' |
|
import numpy as np |
|
import torch |
|
import torch.backends.cudnn as cudnn |
|
from PIL import Image |
|
from models.SwinIR import Generator |
|
from utils.utils import cvtColor, preprocess_input |
|
|
|
|
|
class MASKGAN(object): |
|
|
|
|
|
|
|
_defaults = { |
|
|
|
|
|
|
|
"model_path" : 'model_data/G_FFHQ.pth', |
|
|
|
|
|
|
|
"scale_factor" : 1, |
|
|
|
|
|
|
|
"hr_shape" : [112, 112], |
|
|
|
|
|
|
|
|
|
"cuda" : False, |
|
} |
|
|
|
|
|
|
|
|
|
def __init__(self, **kwargs): |
|
self.__dict__.update(self._defaults) |
|
for name, value in kwargs.items(): |
|
setattr(self, name, value) |
|
self.generate() |
|
|
|
def generate(self): |
|
self.net = Generator(upscale=self.scale_factor, img_size=tuple(self.hr_shape), |
|
window_size=7, img_range=1., depths=[6, 6, 6, 6], |
|
embed_dim=96, num_heads=[6, 6, 6, 6], mlp_ratio=4, upsampler='pixelshuffledirect') |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
self.net = torch.load(self.model_path, map_location=device) |
|
self.net = self.net.eval() |
|
print('{} model, and classes loaded.'.format(self.model_path)) |
|
|
|
if self.cuda: |
|
self.net = torch.nn.DataParallel(self.net) |
|
cudnn.benchmark = True |
|
self.net = self.net.cuda() |
|
|
|
def generate_1x1_image(self, image): |
|
|
|
|
|
|
|
|
|
image = cvtColor(image) |
|
|
|
|
|
|
|
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1]), 0) |
|
|
|
with torch.no_grad(): |
|
image_data = torch.from_numpy(image_data).type(torch.FloatTensor) |
|
if self.cuda: |
|
image_data = image_data.cuda() |
|
|
|
|
|
|
|
|
|
hr_image = self.net(image_data)[0] |
|
|
|
|
|
|
|
hr_image = (hr_image.cpu().data.numpy().transpose(1, 2, 0) * 0.5 + 0.5) |
|
hr_image = np.clip(hr_image * 255, 0, 255) |
|
|
|
hr_image = Image.fromarray(np.uint8(hr_image)) |
|
return hr_image |
|
|
