utils.py

import random
import numpy as np
import xarray as xr
from tqdm import tqdm
import matplotlib.pyplot as plt

def get_phase_and_magnitude(xar_dinsar):
    """
    :param xar_dinsar: dinsar scene. First entry must be real, second must be imaginary
    :return: phase and magnitude
    """
    real = xar_dinsar.values[0, :, :]
    imag = xar_dinsar.values[1, :, :]

    # Create complex
    complex_data = real + 1j * imag
    phase = np.angle(complex_data)
    magnitude = np.abs(complex_data)

    return phase, magnitude


def flow_train_dataset(scene_dinsar, scene_mask, scene_ground_lines, scene_name, patch_size, random_state, path_out, save):

    seen = xr.zeros_like(scene_mask)
    limit = 15000   # upper limit
    no_images_tot = 0   # total img counter
    pbar = tqdm(total=500)

    while True:

        list_patches, list_masks, seen, indexes_x, indexes_y = create_dataset_train(scene_dinsar,
                                                                                    scene_mask,
                                                                                    scene_ground_lines,
                                                                                    patch_size,
                                                                                    seen,
                                                                                    random_state)

        assert len(list_patches) == len(list_masks) == len(indexes_x) == len(indexes_y), "Error occurred."

        no_images_it = len(list_patches)    # iteration count
        no_images_tot += no_images_it       # total count

        if no_images_it == 0:
            print("Unable to find new patches")
            break

        # SAVE THE IMAGES
        for (patch, mask, idx_x, idx_y) in zip(list_patches, list_masks, indexes_x, indexes_y):

            # numpy
            patch_np = patch.values # float32 (2, 128, 128)
            mask_np = mask.values   # uint8 (1, 128, 128)

            if save:
                # out name format
                file_out = f"{path_out}/{patch_np.shape[1]}/{scene_name}_x{idx_x}_y{idx_y}_s{patch_np.shape[1]}"

                # save .npz (re + im + mask)
                np.savez(f"{file_out}.npz", image=patch_np, mask=mask_np)

                # save mask .png
                plt.imsave(f"{file_out}_mask.png", mask_np.squeeze(), cmap="binary")

                # save phase .png
                phase, magnitude = get_phase_and_magnitude(patch)
                phase_norm = (phase + np.pi) / (2 * np.pi) # Normalize phase from [-π, π] → [0, 1]
                plt.imsave(f"{file_out}_phase.png", phase_norm, cmap="hsv")


        pbar.update(no_images_it)

        if no_images_tot >= limit:
            print(f"We have exceeded the limit {limit}")
            break

    pbar.close()

    return seen, no_images_tot


def create_dataset_train(scene_dinsar, scene_mask, scene_ground_lines, patch_size, seen, random_state):

    i_h, i_w = scene_dinsar.shape[1:]
    s_h, s_w = scene_mask.shape[1:]
    p_h, p_w = patch_size, patch_size

    if p_h > i_h:
        raise ValueError(
            "Height of the patch should be less than the height of the image."
        )

    if p_w > i_w:
        raise ValueError(
            "Width of the patch should be less than the width of the image."
        )

    if i_h != s_h:
        raise ValueError(
            "Height of the mask should equal to the height of the image."
        )

    if i_w != s_w:
        raise ValueError(
            "Width of the mask should equal to the width of the image."
            )

    size = p_h * p_w
    patches, masks, idx_x, idx_y = [], [], [], []
    max_iter = 50000
    seed = None if random_state == -1 else random_state
    random.seed(seed)

    for _ in range(max_iter):
        append = True
        # create random indexes
        i_s = random.randint(0, i_h - p_h)
        j_s = random.randint(0, i_w - p_w)
        # calculate the indexes of the center
        center_i = int(i_s + (p_h / 2))
        center_j = int(j_s + (p_w / 2))

        # calculate the patch and the mask patch
        patch = scene_dinsar[:, i_s:i_s + p_h, j_s:j_s + p_w]  # xarray
        mask_patch = scene_mask[:, i_s:i_s + p_h, j_s:j_s + p_w]  # xarray

        # CONDITION 1: presence of grounding line inside the mask - we want to sample grounding line.
        if float(mask_patch.sum()) == 0:
            append = False

        #todo: condition 1 totally excludes noise. But we need to train a model able to predict noise.
        # Therefore we need to generate empty images.

        #print(scene_dinsar.shape, i_s, j_s, patch.shape)
        #print(scene_mask.shape, i_s, j_s, mask_patch.shape)
        #print(float(mask_patch.sum()), append)

        # CONDITION 2: if this region has already been (partially) sampled
        if float(seen[:,center_i-15:center_i+16, center_j-15:center_j+16].sum()) > 200: # discard if the 32x32 region around the center has been previously sampled to some extent.
            append = False

        if append:
            assert mask_patch.size == size, "Something wrong to be checked."
            #seen[:, center_i - 15:center_i + 16, center_j - 15:center_j + 16] += 1  # To keep track of created patch regions, fill a 32x32 box centered on the patch region with 1.
            seen[:, i_s:i_s + p_h, j_s:j_s + p_w] += 1
            patches.append(patch)
            masks.append(mask_patch)
            idx_x.append(i_s)
            idx_y.append(j_s)

            plot = False
            if plot:
                patch_plot = patch.copy(deep=True)
                phase, magnitude = get_phase_and_magnitude(patch_plot)
                patch_plot.values[0,:,:] = phase
                patch_plot.values[1,:,:] = magnitude
                fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(16, 4))

                scene_ground_lines.plot(ax=ax3, ec='cyan', fc='none', lw=2)
                scene_ground_lines.plot(ax=ax4, ec='red', fc='none', lw=2)

                im1 = patch_plot.sel(band=1).plot(ax=ax1, cmap='hsv', vmin=-np.pi, vmax=np.pi,
                                                          add_colorbar=False)
                im2 = mask_patch.plot(ax=ax2, cmap='gray', add_colorbar=False)
                im3 = scene_mask.plot(ax=ax3, cmap='Blues', add_colorbar=False)
                im4 = seen.plot(ax=ax4, cmap='Reds', add_colorbar=False)
                for ax in (ax1, ax2, ax3, ax4):
                    ax.set_aspect('equal')
                ax1.set_title("patch phase")
                ax2.set_title("patch mask")
                ax3.set_title("scene_mask")
                ax4.set_title("seen")

                # Add colorbars below each plot
                cbar1 = plt.colorbar(im1, ax=ax1, orientation='horizontal', pad=0.1, fraction=0.046)
                cbar1.set_label("Phase [rad]")

                cbar2 = plt.colorbar(im2, ax=ax2, orientation='horizontal', pad=0.1, fraction=0.046)
                cbar2.set_label("binary")

                cbar3 = plt.colorbar(im3, ax=ax3, orientation='horizontal', pad=0.1, fraction=0.046)
                cbar3.set_label("binary")

                cbar4 = plt.colorbar(im4, ax=ax4, orientation='horizontal', pad=0.1, fraction=0.046)
                cbar4.set_label("seen occupancy")

                plt.tight_layout()
                plt.show()

        if len(patches) == 10:
            break

    return patches, masks, seen, idx_x, idx_y