基于稀疏编码的图像超分辨率-优快云博客

本文链接：https://blog.youkuaiyun.com/nsjdjdn/article/details/132138936

1.介绍

参考文章：Image Super-Resolution Via Sparse Representation | IEEE Journals & Magazine | IEEE Xplore

2.代码

SR_function.py

import numpy as np

from skimage.transform import resize
from scipy.signal import convolve2d
from scipy import sparse
from tqdm import tqdm


# brojection 模块
def gauss2D(shape,sigma):

    m,n = [(ss-1.)/2. for ss in shape]
    y,x = np.ogrid[-m:m+1,-n:n+1]
    h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
    h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    return h


def backprojection(img_hr, img_lr, maxIter):
    p = gauss2D((5, 5), 1)
    p = np.multiply(p, p)
    p = np.divide(p, np.sum(p))

    for i in range(maxIter):
        img_lr_ds = resize(img_hr, img_lr.shape, anti_aliasing=1)
        img_diff = img_lr - img_lr_ds

        img_diff = resize(img_diff, img_hr.shape)
        img_hr += convolve2d(img_diff, p, 'same')
    return img_hr

#normalize 模块
def normalize_signal(img, img_lr_ori, channel):
    if np.mean(img[:, :, channel]) * 255 > np.mean(img_lr_ori[:, :, channel]):
        ratio = np.mean(img_lr_ori[:, :, channel]) / (np.mean(img[:, :, channel]) * 255)
        img[:, :, channel] = np.multiply(img[:, :, channel], ratio)
    elif np.mean(img[:, :, channel]) * 255 < np.mean(img_lr_ori[:, :, channel]):
        ratio = np.mean(img_lr_ori[:, :, channel]) / (np.mean(img[:, :, channel]) * 255)
        img[:, :, channel] = np.multiply(img[:, :, channel], ratio)
    return img[:, :, channel]

def normalize_max(img):
    for m in range(img.shape[0]):
        for n in range(img.shape[1]):
            if img[m, n, 0] > 1:
                img[m, n, 0] = 1
            if img[m, n, 1] > 1:
                img[m, n, 1] = 1
            if img[m, n, 2] > 1:
                img[m, n, 2] = 1
    return img

#SR模块
def extract_lr_feat(img_lr):
    h, w = img_lr.shape
    img_lr_feat = np.zeros((h, w, 4))

    # 一阶梯度
    hf1 = [[-1, 0, 1], ] * 3
    vf1 = np.transpose(hf1)

    img_lr_feat[:, :, 0] = convolve2d(img_lr, hf1, 'same')
    img_lr_feat[:, :, 1] = convolve2d(img_lr, vf1, 'same')

    #  二阶
    hf2 = [[1, 0, -2, 0, 1], ] * 3
    vf2 = np.transpose(hf2)

    img_lr_feat[:, :, 2] = convolve2d(img_lr, hf2, 'same')
    img_lr_feat[:, :, 3] = convolve2d(img_lr, vf2, 'same')

    return img_lr_feat

def create_list_step(start, stop, step):
    list_step = []
    for i in range(start, stop, step):
        list_step = np.append(list_step, i)
    return list_step

def lin_scale(xh, us_norm):
    hr_norm = np.sqrt(np.sum(np.multiply(xh, xh)))

    if hr_norm > 0:
        s = us_norm * 1.2 / hr_norm
        xh = np.multiply(xh, s)
    return xh

def fss(lmbd, A, b):
    EPS = 1e-9
    x = np.zeros((A.shape[1], 1))    
    grad = np.dot(A, x) + b     
    ma = np.amax(np.multiply(abs(grad), np.isin(x, 0).T), axis=0)
    mi = np.zeros(grad.shape[1])
    for j in range(grad.shape[1]):
        for i in range(grad.shape[0]):
            if grad[i, j] == ma[j]:
                mi[j] = i
                break
    mi = mi.astype(int)   
    while True:

        if np.all(grad[mi]) > lmbd + EPS:
            x[mi] = (lmbd - grad[mi]) / A[mi, mi]
        elif np.all(grad[mi]) < - lmbd - EPS:
            x[mi] = (- lmbd - grad[mi]) / A[mi, mi]
        else:
            if np.all(x == 0):
                break

        while True:
            
            a = np.where(x != 0)
            Aa = A[a, a]
            ba = b[a]
            xa = x[a]

            vect = -lmbd * np.sign(xa) - ba
            x_new = np.linalg.lstsq(Aa, vect)
            idx = np.where(x_new != 0)
            o_new = np.dot((vect[idx] / 2 + ba[idx]).T, x_new[idx]) + lmbd * np.sum(abs(x_new[idx]))

            s = np.where(np.multiply(xa, x_new) < 0) 
            if np.all(s == 0):
                x[a] = x_new
                loss = o_new
                break
            x_min = x_new
            o_min = o_new
            d = x_new - xa
            t = np.divide(d, xa)
            for zd in s.T:
                x_s = xa - d / t[zd]
                x_s[zd] = 0
                idx = np.where(x_s == 0)
                o_s = np.dot((np.dot(Aa[idx, idx], x_s[idx]) / 2 + ba[idx]).T, x_s[idx] + lmbd * np.sum(abs(x_s[idx])))
                if o_s < o_min:
                    x_min = x_s
                    o_min = o_s
            
            x[a] = x_min
            loss = o_min
        
        grad = np.dot(A, sparse.csc_matrix(x)) + b

        ma, mi = max(np.multiply(abs(grad), np.where(x == 0)))
        if ma <= lmbd + EPS:
            break
    
    return x

# 超分
def ScSR(image, size, upsample_factor, Dh, Dl, lmbd, overlap):
    
    patch_size = 3

    img_us = resize(image, size) 
    img_us_height, img_us_width = img_us.shape

    img_hr = np.zeros(img_us.shape)
    cnt_matrix = np.zeros(img_us.shape)

    img_lr_y_feat = extract_lr_feat(img_hr) 
    
    gridx = np.append(create_list_step(0, img_us_width - patch_size - 1, patch_size - overlap), img_us_width - patch_size - 1)
    gridy = np.append(create_list_step(0, img_us_height - patch_size - 1, patch_size - overlap), img_us_height - patch_size - 1)

    count = 0

    for m in tqdm(range(0, len(gridx))):
        for n in range(0, len(gridy)):
            count += 1
            xx = int(gridx[m])
            yy = int(gridy[n])

            us_patch = img_us[yy : yy + patch_size, xx : xx + patch_size]
            us_mean = np.mean(np.ravel(us_patch, order='F'))
            us_patch = np.ravel(us_patch, order='F') - us_mean
            us_norm = np.sqrt(np.sum(np.multiply(us_patch, us_patch)))

            feat_patch = img_lr_y_feat[yy : yy + patch_size, xx : xx + patch_size, :]
            feat_patch = np.ravel(feat_patch, order='F')
            feat_norm = np.sqrt(np.sum(np.multiply(feat_patch, feat_patch)))

            if feat_norm > 1:
                y = np.divide(feat_patch, feat_norm)
            else:
                y = feat_patch

            b = np.dot(np.multiply(Dl.T, -1), y)
            w = fss(lmbd, Dl, b)

            hr_patch = np.dot(Dh, w)
            hr_patch = lin_scale(hr_patch, us_norm)

            hr_patch = np.reshape(hr_patch, (patch_size, -1))
            hr_patch += us_mean

            img_hr[yy : yy + patch_size, xx : xx + patch_size] += hr_patch
            cnt_matrix[yy : yy + patch_size, xx : xx + patch_size] += 1

    index = np.where(cnt_matrix < 1)[0]
    img_hr[index] = img_us[index]

    cnt_matrix[index] = 1
    img_hr = np.divide(img_hr, cnt_matrix)
    
    return img_hr

train_dict.py



import numpy as np
import pickle
from os import listdir
from skimage import io
from skimage.transform import resize, rescale
from scipy.signal import convolve2d
from tqdm import tqdm
from spams import trainDL
def sample_patches(img, patch_size, patch_num, upscale):
    # if img.shape[2] == 3:
    #     hIm = rgb2gray(img)
    # else:
    #     hIm = img
    hIm = img
    # 生成低分辨率对部分
    lIm = rescale(hIm, 1 / upscale)
    lIm = resize(lIm, hIm.shape)
    nrow, ncol = hIm.shape

    x = np.random.permutation(range(nrow - 2 * patch_size)) + patch_size
    y = np.random.permutation(range(ncol - 2 * patch_size)) + patch_size

    X, Y = np.meshgrid(x, y)
    xrow = np.ravel(X, order='F')
    ycol = np.ravel(Y, order='F')

    if patch_num < len(xrow):
        xrow = xrow[0 : patch_num]
        ycol = ycol[0 : patch_num]

    patch_num = len(xrow)

    H = np.zeros((patch_size ** 2, len(xrow)))
    L = np.zeros((4 * patch_size ** 2, len(xrow)))

    # 计算一阶以及二阶梯度
    hf1 = [[-1, 0, 1], ] * 3
    vf1 = np.transpose(hf1)

    lImG11 = convolve2d(lIm, hf1, 'same')
    lImG12 = convolve2d(lIm, vf1, 'same')

    hf2 = [[1, 0, -2, 0, 1], ] * 3
    vf2 = np.transpose(hf2)

    lImG21 = convolve2d(lIm, hf2, 'same')
    lImG22 = convolve2d(lIm, vf2, 'same')

    for i in tqdm(range(patch_num)):
        row = xrow[i]
        col = ycol[i]

        Hpatch = np.ravel(hIm[row : row + patch_size, col : col + patch_size], order='F')
        # Hpatch = np.reshape(Hpatch, (Hpatch.shape[0], 1))
        
        Lpatch1 = np.ravel(lImG11[row : row + patch_size, col : col + patch_size], order='F')
        Lpatch1 = np.reshape(Lpatch1, (Lpatch1.shape[0], 1))
        Lpatch2 = np.ravel(lImG12[row : row + patch_size, col : col + patch_size], order='F')
        Lpatch2 = np.reshape(Lpatch2, (Lpatch2.shape[0], 1))
        Lpatch3 = np.ravel(lImG21[row : row + patch_size, col : col + patch_size], order='F')
        Lpatch3 = np.reshape(Lpatch3, (Lpatch3.shape[0], 1))
        Lpatch4 = np.ravel(lImG22[row : row + patch_size, col : col + patch_size], order='F')
        Lpatch4 = np.reshape(Lpatch4, (Lpatch4.shape[0], 1))

        Lpatch = np.concatenate((Lpatch1, Lpatch2, Lpatch3, Lpatch4), axis=1)
        Lpatch = np.ravel(Lpatch, order='F')

        if i == 0:
            HP = np.zeros((Hpatch.shape[0], 1))
            LP = np.zeros((Lpatch.shape[0], 1))
            # print(HP.shape)
            HP[:, i] = Hpatch - np.mean(Hpatch)
            LP[:, i] = Lpatch
        else:
            HP_temp = Hpatch - np.mean(Hpatch)
            HP_temp = np.reshape(HP_temp, (HP_temp.shape[0], 1))
            HP = np.concatenate((HP, HP_temp), axis=1)
            LP_temp = Lpatch
            LP_temp = np.reshape(LP_temp, (LP_temp.shape[0], 1))
            LP = np.concatenate((LP, LP_temp), axis=1)
    
    return HP, LP

def rnd_smp_patch(img_path, patch_size, num_patch, upsample):
    
    img_dir = listdir(img_path)

    img_num = len(img_dir)
    nper_img = np.zeros((img_num, 1))

    for i in tqdm(range(img_num)):
        img = io.imread('{}{}'.format(img_path, img_dir[i]))
        nper_img[i] = img.shape[0] * img.shape[1]

    nper_img = np.floor(nper_img * num_patch / np.sum(nper_img, axis=0))

    for i in tqdm(range(img_num)):
        patch_num = int(nper_img[i])
        img = io.imread('{}{}'.format(img_path, img_dir[i]))
        H, L = sample_patches(img, patch_size, patch_num, upsample)
        if i == 0:
            Xh = H
            Xl = L
        else:
            Xh = np.concatenate((Xh, H), axis=1)
            Xl = np.concatenate((Xl, L), axis=1)
    return Xh, Xl

def patch_pruning(Xh, Xl):
    pvars = np.var(Xh, axis=0)
    threshold = np.percentile(pvars, 10)
    idx = pvars > threshold
    # print(pvars)
    Xh = Xh[:, idx]
    Xl = Xl[:, idx]
    return Xh, Xl

# 参数设置

dict_size = 1024
lmbd = 0.1
patch_size = 3
num_samples = 1000000
upsample_factor = 2

# 读取路径
train_img_path = r'D:\pycharm\pytorch\study\data\sonar_ship/'

Xh, Xl = rnd_smp_patch(train_img_path, patch_size, num_samples, upsample_factor)

Xh, Xl = patch_pruning(Xh, Xl)
Xh = np.asfortranarray(Xh)
Xl = np.asfortranarray(Xl)

# 字典学习
Dh = trainDL(Xh, K=dict_size, lambda1=lmbd, iter=100)
Dl = trainDL(Xl, K=dict_size, lambda1=lmbd, iter=100)

# 保存路径
with open('data/dicts/'+ 'Dh_' + str(dict_size) + '_US' + str(upsample_factor) + '_L' + str(lmbd) + '_PS' + str(patch_size) + '.pkl', 'wb') as f:
    pickle.dump(Dh, f, pickle.HIGHEST_PROTOCOL)

with open('data/dicts/'+ 'Dl_' + str(dict_size) + '_US' + str(upsample_factor) + '_L' + str(lmbd) + '_PS' + str(patch_size) + '.pkl', 'wb') as f:
    pickle.dump(Dl, f, pickle.HIGHEST_PROTOCOL)

makedata_bicubic.py


from os import listdir
import cv2
from tqdm import tqdm
###需要修改的部分
train_hr_path = 'path'
train_lr_path = 'path'

val_hr_path = 'path'
val_lr_path = 'path'

scale = 2
###

# 训练集的下采样
num_train_images = len(listdir(train_hr_path))
for i in tqdm(range(num_train_images)):
    img_name = listdir(train_hr_path)[i]
    img = cv2.imread('{}{}'.format(train_hr_path, img_name),0)
    scaled_image = cv2.resize(img, (0, 0), fx=(1/scale), fy=(1/scale), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite('{}{}'.format(train_lr_path, img_name), scaled_image)

# 验证集的下采样

num_val_images = len(listdir(val_hr_path))
for i in tqdm(range(num_val_images)):
    img_name = listdir(val_hr_path)[i]
    img = cv2.imread('{}{}'.format(val_hr_path, img_name),0)
    cv2.imwrite('{}{}'.format(val_hr_path, img_name), img)
    scaled_image = cv2.resize(img, (0, 0), fx=(1/scale), fy=(1/scale), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite('{}{}'.format(val_lr_path, img_name), scaled_image)

val_image.py

import cv2
import numpy as np
import pickle

from os import listdir, makedirs
from os.path import isdir

from skimage.color import rgb2ycbcr, ycbcr2rgb, rgb2gray
from skimage.transform import resize

from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import normalize

from tqdm import tqdm




from ScSR import ScSR
from backprojection import backprojection

def normalize_signal(img, img_lr_ori, channel):
    if np.mean(img[:, :, channel]) * 255 > np.mean(img_lr_ori[:, :, channel]):
        ratio = np.mean(img_lr_ori[:, :, channel]) / (np.mean(img[:, :, channel]) * 255)
        img[:, :, channel] = np.multiply(img[:, :, channel], ratio)
    elif np.mean(img[:, :, channel]) * 255 < np.mean(img_lr_ori[:, :, channel]):
        ratio = np.mean(img_lr_ori[:, :, channel]) / (np.mean(img[:, :, channel]) * 255)
        img[:, :, channel] = np.multiply(img[:, :, channel], ratio)
    return img[:, :, channel]

def normalize_max(img):
    for m in range(img.shape[0]):
        for n in range(img.shape[1]):
            if img[m, n, 0] > 1:
                img[m, n, 0] = 1
            if img[m, n, 1] > 1:
                img[m, n, 1] = 1
            if img[m, n, 2] > 1:
                img[m, n, 2] = 1
    return img
###需要修改的部分
D_size = 1024
US_mag = 2
lmbd = 0.1
patch_size = 3
###需要修改的部分

dict_name = str(D_size) + '_US' + str(US_mag) + '_L' + str(lmbd) + '_PS' + str(patch_size)

with open('data/dicts/Dh_' + dict_name + '.pkl', 'rb') as f:
    Dh = pickle.load(f)
Dh = normalize(Dh)
with open('data/dicts/Dl_' + dict_name + '.pkl', 'rb') as f:
    Dl = pickle.load(f)
Dl = normalize(Dl)

img_lr_dir = 'data/val_lr/'
img_hr_dir = 'data/val_hr/'
###可以修改的部分
overlap = 1
lmbd = 0.3
upsample = 2
max_iteration = 1000
###


img_lr_file = listdir(img_lr_dir)

for i in tqdm(range(len(img_lr_file))):
    # for i in tqdm(range(1)):#单张验证
    img_name = img_lr_file[i]
    img_name_dir = list(img_name)
    img_name_dir = np.delete(np.delete(np.delete(np.delete(img_name_dir, -1), -1), -1), -1)
    img_name_dir = ''.join(img_name_dir)
    if isdir('path' + dict_name + '_' + img_name_dir) == False:
        new_dir = makedirs('{}{}'.format('path' + dict_name + '_', img_name_dir))
    img_lr = cv2.imread('{}{}'.format(img_lr_dir, img_name))

    ## 读取和保存图片
    img_hr = cv2.imread('{}{}'.format(img_hr_dir, img_name))
    cv2.imwrite('{}{}{}{}'.format('data/results/set5_sigma25/' + dict_name + '_', img_name_dir, '/', '3GT.png'), img_hr)

    # 改变颜色空间
    img_hr_y = rgb2ycbcr(img_hr)[: ,: ,0]
    img_lr_ori = img_lr
    temp = img_lr
    imr_lr = rgb2ycbcr(img_lr)
    img_lr_y = img_lr[: ,: ,0]
    img_lr_cb = img_lr[: ,: ,1]
    img_lr_cr = img_lr[: ,: ,2]

    img_sr_cb = resize(img_lr_cb, (img_hr.shape[0], img_hr.shape[1]), order=0)
    img_sr_cr = resize(img_lr_cr, (img_hr.shape[0], img_hr.shape[1]), order=0)

    # 超分
    img_sr_y = ScSR(img_lr_y, img_hr_y.shape, upsample, Dh, Dl, lmbd, overlap)
    img_sr_y = backprojection(img_sr_y, img_lr_y, max_iteration)
    # img_sr_y = resize(img_lr_y, (img_hr.shape[0], img_hr.shape[1]), order=0) # Loop check
    img_sr = np.stack((img_sr_y, img_sr_cb, img_sr_cr), axis=2)
    img_sr = ycbcr2rgb(img_sr)

    for channel in range(img_sr.shape[2]):
        img_sr[:, :, channel] = normalize_signal(img_sr, img_lr_ori, channel)

    img_sr = normalize_max(img_sr)



    ## pnsr的计算
    rmse_sr_hr = np.sqrt(mean_squared_error(img_hr_y, img_sr_y))
    # psnr_sr_hr = 10*np.log10(255.0**2/rmse_sr_hr**2)
    psnr_sr_hr = 10 *np.log10(1.0** 2 /rmse_sr_hr**2)
    psnr_sr_hr = np.zeros((1,)) + psnr_sr_hr
    np.savetxt('{}{}{}{}'.format('path' + dict_name + '_', img_name_dir, '/', 'PSNR_SR.txt'), psnr_sr_hr)

    ## 保存图片
    cv2.imwrite('{}{}{}{}'.format('path' + dict_name + '_', img_name_dir, '/', '2SR.png'), img_sr)