MRI_recon/code/Frequency-Diffusion/draw/frequency_sampling.py

import torch

from utils.k_degrade_utils import *


if __name__ == "__main__":
    # First STEP
    import matplotlib.pyplot as plt
    import numpy as np, os

    os.makedirs("outputs", exist_ok=True)

    os.makedirs("outputs/low-fre-first", exist_ok=True)
    os.makedirs("outputs/random-sample", exist_ok=True)
    
    
    image_size = 256
    accelerated_factor = 6
    center_fraction = 0.04
    time_step = 25


    masks = get_ksu_kernel(time_step, image_size, "LogSamplingRate",
                           accelerated_factor=accelerated_factor, center_fraction=center_fraction) # LogSamplingRate


    batch_size = 1

    img = plt.imread("./assets/BraTS20_Training_001_86_t1.png")
    img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
    img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)

    print("input img shape: ", img.shape)

    # to gray scale
    if len(img.shape) == 3 and img.shape[-1] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)


    # img = np.transpose(img, (2, 0, 1))
    # img = img[0]
    img = np.expand_dims(img, axis=0)
    img = torch.from_numpy(img).unsqueeze(0).float()
    original_img = img.clone()


    rand_kernels = []
    rand_x = torch.randint(0, image_size + 1, (batch_size,)).long()

    img = img #* 2 - 1  #

    masked_img = []

    for m in masks:
        m = m.unsqueeze(0)
        img = apply_ksu_kernel(img, m)
        masked_img.append(img)

    save_masks = masks
    masks = np.concatenate(masks, axis=-1)[0]
    masked_img = torch.concat(masked_img, dim=-1).numpy()  #+ 1) * 0.5

    masked_img = np.transpose(masked_img, (0, 2, 3, 1))[0, ..., 0]
    # masked_img = cv2.cvtColor(masked_img, cv2.COLOR_RGB2GRAY)



    img = np.concatenate([masks, masked_img], axis=0)
    min_ = masked_img.min()
    max_ = masked_img.max()

    out = img[image_size: 2 * image_size, :  image_size]
    fft, _ = apply_tofre(torch.from_numpy(out), torch.from_numpy(out))  # complex
    fft = np.abs(fft.numpy())
    fft = np.log(fft)
    fft = (fft - fft.min()) / (fft.max() - fft.min())

    for i in range(time_step+1):
        out = img[image_size: 2 * image_size, i * image_size: (i + 1) * image_size]

        # out = (out - out.min()) / (out.max() - out.min())
        out = (out - min_) / (max_ - min_)
        plt.imsave(f"outputs/low-fre-first/{i}_image.png", out, cmap='gray')

        if i != 0:
            out = img[:image_size, i * image_size:(i + 1) * image_size]
            out = (out - out.min()) / (out.max() - out.min())

            plt.imsave(f"outputs/low-fre-first/{i}_mask.png", out, cmap='gray')

            save_fft = fft * out
            plt.imsave(f"outputs/low-fre-first/{i}_fft.png", save_fft, cmap='gray')


        else:
            diff = np.ones((image_size, image_size, 3), dtype=np.uint8) * 255  # All 255 (White)ve
            ones = diff.astype(np.float32) / 255.0
            print("ones shape: ", ones.shape, ones.min(), ones.max())

            plt.imsave(f"outputs/low-fre-first/{i}_mask.png", ones, cmap='gray')
            plt.imsave(f"outputs/low-fre-first/{i}_fft.png", fft, cmap='gray')

        try:
            diff = img[:image_size, (i-1) * image_size:(i) * image_size] - \
                   img[:image_size, (i) * image_size:(i + 1) * image_size]

        except:
            diff = np.zeros_like(img[:image_size, : image_size])

        # plt.imsave(f"outputs/low-fre-first/{i}_mask_diff.png", diff, cmap='gray')
        # print("diff shape: ", diff.shape, diff.min(), diff.max())


        diffsig = diff * fft
        # save it as a red img, but the bg is trasparent
        alpha_channel = np.full_like(diff, 255, dtype=np.uint8) * diff
        alpha_channel = np.expand_dims(alpha_channel, axis=-1)

        diff = (diff * 255).astype(np.uint8)
        diff = np.stack([diff, np.zeros_like(diff), np.zeros_like(diff)], axis=-1)
        # Create an alpha channel (255 for full opacity)

        # Concatenate RGB with Alpha channel
        diff = np.concatenate([diff, alpha_channel], axis=-1)
        diff = diff.astype(np.uint8)

        # print("diff shape: ", diff.shape, diff.min(), diff.max())


        plt.imsave(f"outputs/low-fre-first/{i}_mask_diff_red.png", diff, cmap='gray')
        plt.imsave(f"outputs/low-fre-first/{i}_mask_diffsig.png", diffsig, cmap='gray')
        

    plt.imsave("outputs/masked_img.png", masked_img, cmap='gray')
    plt.figure(figsize=(5*time_step, 10))
    plt.imshow(img, cmap='gray')      # (1, 128, 1280)
    plt.show()

    print("\n\nSecond stage...")

    # -------------------------------   -------------------------------   -------------------------------
    # -------------------------------   -------------------------------   -------------------------------

    # Second STEP  completely Random
    import matplotlib.pyplot as plt
    import numpy as np


    final_mask = save_masks[-1][0].numpy()
    new_masks = []

    plt.imshow(final_mask, cmap='gray')
    plt.show()

    height, width = final_mask.shape
    print("final_mask shape: ", final_mask.shape)

    # Count ones and zeros
    ones  = np.sum(final_mask[0] == 1)
    zeros = np.sum(final_mask[0] == 0)

    print("Initial ones count:", ones)
    print("Initial zeros count:", zeros)

    # Identify initially filled and empty strips
    initial_filled_indices = np.where(final_mask[0] == 1)[0]
    remaining_indices = np.where(final_mask[0] == 0)[0]

    # Shuffle remaining indices to randomize filling order
    np.random.shuffle(remaining_indices)

    # Split remaining indices into `time_step` parts
    fills_per_step = np.array_split(remaining_indices, time_step)

    masked_img = []  # Store masks at each step

    # Copy initial mask
    current_mask = final_mask.copy()
    new_masks.append(current_mask.copy())  # Store initial state

    # Fill remaining strips over time
    for i in range(time_step):
        current_mask[:, fills_per_step[i - 1]] = 1  # Fill new strips
        new_masks.append(current_mask.copy())  # Store new mask
    # current_mask.append(final_mask)  # Store new mask

    new_masks = new_masks[::-1]  # Reverse list to get correct order
    masked_img = []

    for m in new_masks:
        m = torch.from_numpy(m)  #.unsqueeze(0)

        img = apply_ksu_kernel(original_img, m)
        masked_img.append(img)

    masks = np.concatenate(new_masks, axis=-1)
    masked_img = torch.concat(masked_img, dim=-1).numpy()  #+ 1) * 0.5

    masked_img = np.transpose(masked_img, (0, 2, 3, 1))[0, ..., 0]

    print("masked_img shape: ", masked_img.shape)


    # masked_img = cv2.cvtColor(masked_img, cv2.COLOR_RGB2GRAY)
    # masked_img = (masked_img - masked_img.min()) / (masked_img.max() - masked_img.min())

    img = np.concatenate([masks, masked_img], axis=0)
    
    
    min_ = masked_img.min()
    max_ = masked_img.max()

    out = img[image_size: 2 * image_size, :  image_size]
    fft, _ = apply_tofre(torch.from_numpy(out), torch.from_numpy(out))  # complex
    fft = np.abs(fft.numpy())
    fft = np.log(fft)
    fft = (fft - fft.min()) / (fft.max() - fft.min())


    for i in range(time_step+1):
        # if i % 3 != 0:
        #     continue
        out = img[image_size : 2*image_size, i * image_size : (i + 1) * image_size]

        # out = (out - out.min()) / (out.max() - out.min())
        out = (out - min_) / (max_ - min_)
        plt.imsave(f"outputs/random-sample/{i}_image.png", out, cmap='gray')

        if i != 0:
            out = img[:image_size, i * image_size:(i + 1) * image_size]
            out = (out - out.min()) / (out.max() - out.min())

            plt.imsave(f"outputs/random-sample/{i}_mask.png", out, cmap='gray')

            save_fft = fft * out
            plt.imsave(f"outputs/random-sample/{i}_fft.png", save_fft, cmap='gray')

            noise = np.random.normal(0, 0.2*np.log((time_step-i)+1), out.shape) * fft
            save_fft = fft + noise * (1-out) # Sigma
            plt.imsave(f"outputs/random-sample/{i}_fft_reverse.png", save_fft, cmap='gray')


        else:
            ones = np.ones_like(out) * 255
            plt.imsave(f"outputs/random-sample/{i}_mask.png", ones, cmap='gray')
            plt.imsave(f"outputs/random-sample/{i}_fft.png", fft, cmap='gray')


        try:
            diff = img[:image_size, (i-1) * image_size:(i) * image_size] - \
                   img[:image_size, (i) * image_size:(i + 1) * image_size]

        except:
            diff = np.zeros_like(img[:image_size, : image_size])


        plt.imsave(f"outputs/random-sample/{i}_mask_diff.png", diff, cmap='gray')
        # print("diff shape: ", diff.shape, diff.min(), diff.max())

        # save it as a red img, but the bg is trasparent
        alpha_channel = np.full_like(diff, 255, dtype=np.uint8) * diff
        alpha_channel = np.expand_dims(alpha_channel, axis=-1)

        diff = (diff * 255).astype(np.uint8)
        diff = np.stack([diff, np.zeros_like(diff), np.zeros_like(diff)], axis=-1)
        # Create an alpha channel (255 for full opacity)

        # Concatenate RGB with Alpha channel
        diff = np.concatenate([diff, alpha_channel], axis=-1)
        diff = diff.astype(np.uint8)

        print("diff shape: ", diff.shape, diff.min(), diff.max())

        plt.imsave(f"outputs/random-sample/{i}_mask_diff_red.png", diff, cmap='gray')


    plt.imsave("outputs/img.png", img, cmap='gray')
    # plt.figure(figsize=(5*time_step, 10))

    plt.imshow(img, cmap='gray')      # (1, 128, 1280)
    plt.tight_layout()
    plt.show()

    print("\n\nSecond stage...")