Example #1
0
def inpainting_lista(missing_pixels, dicsize, n_nonzero_coefs, blksize,
                     overlap, maxiter, grayimg):
    img_ori, masked_img, patchvec, maskvec = generate_inpainting_data(
        missing_pixels=missing_pixels,
        blksize=blksize,
        overlap=overlap,
        grayimg=grayimg)

    # settings
    h, w = img_ori.shape[0], img_ori.shape[1]
    m = patchvec.shape[0]
    n = dicsize
    s = n_nonzero_coefs / n

    print('[INFO] Dictionary size ({}, {}).'.format(m, n))

    print('[INFO] Sparsity {}.'.format(n_nonzero_coefs))
    print('[INFO] OMP max iter {}.'.format(maxiter))

    D = get_init_dict(dicsize=dicsize, blksize=blksize)
    Y = patchvec

    # training
    model = train_lista(D, Y, maskvec, s, img_ori, blksize, overlap)

    # testing
    with torch.no_grad():
        Y_pt = torch.from_numpy(Y.T / 255.0).cuda()
        coef_pt = model(Y_pt)
        coef = coef_pt.cpu().detach().numpy().T

    # testing
    recovered_img = patch2im_for_inpainting(patch_vecs=255.0 *
                                            np.matmul(D, coef),
                                            mask_vecs=maskvec,
                                            imgsize=img_ori.shape,
                                            blksize=blksize,
                                            overlap=overlap)

    # save results
    if not os.path.exists('./save_lista'):
        os.mkdir('./save_lista')

    np.save('./save_lista/img_ori', img_ori)
    np.save('./save_lista/masked_img', masked_img)
    np.save('./save_lista/recovered_img', recovered_img)

    return img_ori, masked_img, recovered_img
def denoising_omp(dicsize, sparsity, blksize, overlap, maxiter):

    np.random.seed(0)

    img_ori, masked_img, mask, patchvec, maskvec = generate_denoising_data(blksize, overlap)

    # settings
    m = patchvec.shape[0]
    n = dicsize

    print('[INFO] Dictionary size ({}, {}).'.format(m, n))
    print('[INFO] Sparsity {}.'.format(sparsity))
    print('[INFO] De-noising max iter {}.'.format(maxiter))

    load_results = False
    if load_results:
        print('[INFO] Load OMP results.')
        dict = np.load('./dict.npy')
        coef = np.load('./coef.npy')

        # reconstruct image
        print('[INFO] Reconstruction ...')
        recovered_img = patch2im_for_denoising(image=img_ori, patch_vecs=np.matmul(dict, coef), imgsize=img_ori.shape,
                                               blksize=blksize, overlap=overlap)

        return img_ori, masked_img, recovered_img

    else:
        # OMP for de-noising
        print('[INFO] Run OMP for denoising')

        recovered_img = None
        for iter in range(maxiter):
            D = get_init_dict(dicsize=dicsize, blksize=blksize)
            dict, coef = denoise_omp_once(patchvec, D, sparsity)

            np.save('./dict', dict)
            np.save('./coef', coef)

            # reconstruct image
            print('[INFO] Reconstruction ...')
            recovered_img = patch2im_for_denoising(image=img_ori, patch_vecs=np.matmul(dict, coef), imgsize=img_ori.shape,
                                                   blksize=blksize, overlap=overlap)

        return img_ori, masked_img, recovered_img
Example #3
0
def inpainting_ksvd(missing_pixels, dicsize, n_nonzero_coefs, blksize, overlap, maxiter, grayimg=False):
    img_ori, masked_img, patchvec, maskvec = generate_inpainting_data(missing_pixels=missing_pixels,
                                                                      blksize=blksize, overlap=overlap, grayimg=grayimg)
    # img_ori, masked_img, patchvec, maskvec = load_inpainting_image(blksize=blksize, overlap=overlap)

    # settings
    h, w = img_ori.shape[0], img_ori.shape[1]
    m = patchvec.shape[0]
    n = dicsize

    print('[INFO] Dictionary size ({}, {}).'.format(m, n))

    print('[INFO] Sparsity {}.'.format(n_nonzero_coefs))
    print('[INFO] OMP max iter {}.'.format(maxiter))

    load_results = True
    if load_results:
        print('[INFO] Load OMP results.')
        if grayimg:
            dict = np.load('./D.npy')
            coef = np.load('./coef.npy')
            recovered_img = patch2im_for_inpainting(patch_vecs=np.matmul(dict, coef), mask_vecs=maskvec,
                                                    imgsize=(h, w), blksize=blksize, overlap=overlap)
        else:
            dictR = np.load('./DR.npy')
            dictG = np.load('./DG.npy')
            dictB = np.load('./DB.npy')
            coefR = np.load('./coefR.npy')
            coefG = np.load('./coefG.npy')
            coefB = np.load('./coefB.npy')

            recovered_imgR = patch2im_for_inpainting(patch_vecs=np.matmul(dictR, coefR), mask_vecs=maskvec,
                                                     imgsize=(h, w), blksize=blksize, overlap=overlap)
            recovered_imgG = patch2im_for_inpainting(patch_vecs=np.matmul(dictG, coefG), mask_vecs=maskvec,
                                                     imgsize=(h, w), blksize=blksize, overlap=overlap)
            recovered_imgB = patch2im_for_inpainting(patch_vecs=np.matmul(dictB, coefB), mask_vecs=maskvec,
                                                     imgsize=(h, w), blksize=blksize, overlap=overlap)
            recovered_img = np.stack((recovered_imgR, recovered_imgG, recovered_imgB), axis=2)

        return img_ori, masked_img, recovered_img

    else:
        # OMP for inpainting
        print('[INFO] Run OMP for image inpainting')

        if grayimg:
            D = get_init_dict(dicsize=dicsize, blksize=blksize)
            Y = patchvec
            newD, coef = inpainting_ksvd_once(Y, D, maskvec, n_nonzero_coefs=50, maxiter=maxiter)

            print('[INFO] Saving dictionary ...')
            np.save('./D', newD)

            print('[INFO] Saving coefficients ...')
            np.save('./coef', coef)

            # recover data
            print('[INFO] Reconstruction ...')

            recovered_img = patch2im_for_inpainting(patch_vecs=np.matmul(newD, coef), mask_vecs=maskvec, imgsize=img_ori.shape,
                                                    blksize=blksize, overlap=overlap)
        else:
            D = get_init_dict(dicsize=dicsize, blksize=blksize)
            RGBname = ['R', 'G', 'B']
            recovered_img = np.zeros_like(img_ori)
            # do R-G-B
            for channel in range(3):
                Y = patchvec[:, :, channel]
                newD, coef = inpainting_ksvd_once(Y, D, maskvec, n_nonzero_coefs=50, maxiter=maxiter)

                print('[INFO] Saving dictionary ...')
                np.save('./D{}'.format(RGBname[channel]), newD)

                print('[INFO] Saving coefficients ...')
                np.save('./coef{}'.format(RGBname[channel]), coef)

                # recover data
                print('[INFO] Reconstruction ...')

                recovered_img_c = patch2im_for_inpainting(patch_vecs=np.matmul(newD, coef), mask_vecs=maskvec, imgsize=img_ori.shape,
                                                          blksize=blksize, overlap=overlap)
                recovered_img[:, :, channel] = recovered_img_c

    return img_ori, masked_img, recovered_img