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...")