修复

utils/dataloader.py

@@ -1,101 +0,0 @@
-from random import randint
-
-import cv2
-import numpy as np
-from PIL import Image
-from torch.utils.data.dataset import Dataset
-
-from .utils import cvtColor, preprocess_input
-
-def look_image(image_name, image):
-    image = np.array(image)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    cv2.imshow(image_name, image)
-    cv2.waitKey(0)
-
-
-def get_new_img_size(width, height, img_min_side=600):
-    if width <= height:
-        f = float(img_min_side) / width
-        resized_height = int(f * height)
-        resized_width = int(img_min_side)
-    else:
-        f = float(img_min_side) / height
-        resized_width = int(f * width)
-        resized_height = int(img_min_side)
-
-    return resized_width, resized_height
-
-class MASKGANDataset(Dataset):
-    def __init__(self, train_lines, lr_shape, hr_shape):
-        super(MASKGANDataset, self).__init__()
-
-        self.train_lines    = train_lines
-        self.train_batches  = len(train_lines)
-
-        self.lr_shape       = lr_shape
-        self.hr_shape       = hr_shape
-
-    def __len__(self):
-        return self.train_batches
-
-    def __getitem__(self, index):
-        index = index % self.train_batches
-        image_list = self.train_lines[index].split(' ')
-        image_origin = Image.open(image_list[0])
-        image_masked = Image.open(image_list[1].split()[0])
-
-        image_origin, image_masked = self.get_random_data(image_origin, image_masked, self.hr_shape)
-
-        image_origin = image_origin.resize((self.hr_shape[1], self.hr_shape[0]), Image.BICUBIC)
-        image_masked = image_masked.resize((self.lr_shape[1], self.lr_shape[0]), Image.BICUBIC)
-        # look_image('origin', image_origin)
-        # look_image('masked', image_masked)
-        image_origin = np.transpose(preprocess_input(np.array(image_origin, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
-        image_masked = np.transpose(preprocess_input(np.array(image_masked, dtype=np.float32), [0.5,0.5,0.5], [0.5,0.5,0.5]), [2,0,1])
-
-        return np.array(image_masked), np.array(image_origin)
-
-    def rand(self, a=0, b=1):
-        return np.random.rand()*(b-a) + a
-
-    def get_random_data(self, image_origin, image_masked, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5, random=True):
-        #------------------------------#
-        #   读取图像并转换成RGB图像
-        #------------------------------#
-        image_origin   = cvtColor(image_origin)
-        image_masked   = cvtColor(image_masked)
-
-        #------------------------------------------#
-        #   色域扭曲
-        #------------------------------------------#
-        hue = self.rand(-hue, hue)
-        sat = self.rand(1, sat) if self.rand()<.5 else 1/self.rand(1, sat)
-        val = self.rand(1, val) if self.rand()<.5 else 1/self.rand(1, val)
-
-        x = cv2.cvtColor(np.array(image_origin,np.float32)/255, cv2.COLOR_RGB2HSV)
-        x[..., 1] *= sat
-        x[..., 2] *= val
-        x[x[:,:, 0]>360, 0] = 360
-        x[:, :, 1:][x[:, :, 1:]>1] = 1
-        x[x<0] = 0
-        image_data_origin = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
-
-        x = cv2.cvtColor(np.array(image_masked,np.float32)/255, cv2.COLOR_RGB2HSV)
-        x[..., 1] *= sat
-        x[..., 2] *= val
-        x[x[:,:, 0]>360, 0] = 360
-        x[:, :, 1:][x[:, :, 1:]>1] = 1
-        x[x<0] = 0
-        image_data_masked = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
-
-        return Image.fromarray(np.uint8(image_data_origin)), Image.fromarray(np.uint8(image_data_masked))
-
-        
-def MASKGAN_dataset_collate(batch):
-    images_l = []
-    images_h = []
-    for img_l, img_h in batch:
-        images_l.append(img_l)
-        images_h.append(img_h)
-    return np.array(images_l), np.array(images_h)

utils/utils.py

@@ -3,7 +3,6 @@ import numpy as np
 import matplotlib.pyplot as plt
 import torch
 from torch.nn import functional as F
-import cv2
 import distutils.util
 
 def show_result(num_epoch, G_net, imgs_lr, imgs_hr):
@@ -63,102 +62,3 @@ def add_arguments(argname, type, default, help, argparser, **kwargs):
                            help=help + ' 默认: %(default)s.',
                            **kwargs)
 
-def filter2D(img, kernel):
-    """PyTorch version of cv2.filter2D
-
-    Args:
-        img (Tensor): (b, c, h, w)
-        kernel (Tensor): (b, k, k)
-    """
-    k = kernel.size(-1)
-    b, c, h, w = img.size()
-    if k % 2 == 1:
-        img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
-    else:
-        raise ValueError('Wrong kernel size')
-
-    ph, pw = img.size()[-2:]
-
-    if kernel.size(0) == 1:
-        # apply the same kernel to all batch images
-        img = img.view(b * c, 1, ph, pw)
-        kernel = kernel.view(1, 1, k, k)
-        return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
-    else:
-        img = img.view(1, b * c, ph, pw)
-        kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
-        return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
-
-
-def usm_sharp(img, weight=0.5, radius=50, threshold=10):
-    """USM sharpening.
-
-    Input image: I; Blurry image: B.
-    1. sharp = I + weight * (I - B)
-    2. Mask = 1 if abs(I - B) > threshold, else: 0
-    3. Blur mask:
-    4. Out = Mask * sharp + (1 - Mask) * I
-
-
-    Args:
-        img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
-        weight (float): Sharp weight. Default: 1.
-        radius (float): Kernel size of Gaussian blur. Default: 50.
-        threshold (int):
-    """
-    if radius % 2 == 0:
-        radius += 1
-    blur = cv2.GaussianBlur(img, (radius, radius), 0)
-    residual = img - blur
-    mask = np.abs(residual) * 255 > threshold
-    mask = mask.astype('float32')
-    soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
-
-    sharp = img + weight * residual
-    sharp = np.clip(sharp, 0, 1)
-    return soft_mask * sharp + (1 - soft_mask) * img
-
-
-class USMSharp(torch.nn.Module):
-
-    def __init__(self, radius=50, sigma=0):
-        super(USMSharp, self).__init__()
-        if radius % 2 == 0:
-            radius += 1
-        self.radius = radius
-        kernel = cv2.getGaussianKernel(radius, sigma)
-        kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0)
-        self.register_buffer('kernel', kernel)
-
-    def forward(self, img, weight=0.5, threshold=10):
-        blur = filter2D(img, self.kernel)
-        residual = img - blur
-
-        mask = torch.abs(residual) * 255 > threshold
-        mask = mask.float()
-        soft_mask = filter2D(mask, self.kernel)
-        sharp = img + weight * residual
-        sharp = torch.clip(sharp, 0, 1)
-        return soft_mask * sharp + (1 - soft_mask) * img
-
-class USMSharp_npy():
-
-    def __init__(self, radius=50, sigma=0):
-        super(USMSharp_npy, self).__init__()
-        if radius % 2 == 0:
-            radius += 1
-        self.radius = radius
-        kernel = cv2.getGaussianKernel(radius, sigma)
-        self.kernel = np.dot(kernel, kernel.transpose()).astype(np.float32)
-
-    def filt(self, img, weight=0.5, threshold=10):
-        blur = cv2.filter2D(img, -1, self.kernel)
-        residual = img - blur
-
-        mask = np.abs(residual) * 255 > threshold
-        mask = mask.astype(np.float32)
-        soft_mask = cv2.filter2D(mask, -1, self.kernel)
-        sharp = img + weight * residual
-        sharp = np.clip(sharp, 0, 1)
-        return soft_mask * sharp + (1 - soft_mask) * img
-