utils/utilities.py

import h5py
import numpy as np
import torch


def compute_proba_batch(weights, logits, mask=None, eps=1e-8):
    if mask is None:
        normalized_weights = torch.softmax(weights, dim=1)
    else:
        normalized_weights = weights.masked_fill(~mask, float("-inf"))
        normalized_weights = torch.softmax(normalized_weights, dim=1)
        normalized_weights = torch.nan_to_num(normalized_weights, nan=0.0)

    weighted = normalized_weights.unsqueeze(2) * logits
    numerator = weighted.sum(dim=1)
    return numerator

# Mapping original → group_id
orig_to_group = {
    0: 0, 1: 0, 2: 0, 3: 0, 4: 0,      # Pylon
    5: 1,                              # Conductor cable
    6: 2, 7: 2,                        # Structural cable
    8: 3, 9: 3,10: 3,11: 3,            # Insulator
    14:4,                              # High vegetation
    15:5,                              # Low vegetation
    16:6,                              # Herbaceous vegetation
    17:7,18:7,                         # Rock, gravel, soil
    19:8,                              # Impervious soil (Road)
    20:9,                              # Water
    21:10,                             # Building
    12:255,13:255,255:255              # Unassigned/Unlabeled
}

_label_lookup = np.full(256, 255, dtype=np.int16)
for _src, _dst in orig_to_group.items():
    if 0 <= int(_src) < _label_lookup.shape[0]:
        _label_lookup[int(_src)] = int(_dst)

# Group names used for display (0–10)
display_names = [
    "Pylon",
    "Conductor cable",
    "Structural cable",
    "Insulator",
    "High vegetation",
    "Low vegetation",
    "Herbaceous vegetation",
    "Rock, gravel, soil",
    "Impervious soil (Road)",
    "Water",
    "Building"
]

def get_main_class_index(orig_label):
    grp = orig_to_group.get(int(orig_label), 255)
    return None if grp == 255 else grp


def map_main_class_indices(labels):
    labels = np.asarray(labels, dtype=np.int64)
    clipped = np.where((labels >= 0) & (labels < _label_lookup.shape[0]), labels, 255)
    return _label_lookup[clipped]

def sample_indices(points, n_per_class, random_state=None):
    """
    For each group 0..10, sample up to ``n_per_class`` indices
    from ``points["GT"]`` without duplicates.
    
    Parameters:
      points: dict containing "GT" (np.array of shape (N,) with original labels)
      n_per_class: int, maximum number of indices to sample per group
      random_state: random seed for reproducibility
    
    Returns:
      np.array of distinct sampled indices, shuffled.
    """
    print("Starting sampling")
    rng = np.random.default_rng(random_state)
    gt = points["GT"]
    N = gt.shape[0]
    n_groups = len(display_names)  # 11

    # One bucket per group.
    selected = {g: [] for g in range(n_groups)}

    # Random permutation of all indices.
    perm = rng.permutation(N)
    

    # Iterate through the permutation and fill each bucket up to n_per_class.
    for idx in perm:
        grp = get_main_class_index(gt[idx])
        if grp is None:
            continue
        if len(selected[grp]) < n_per_class:
            selected[grp].append(idx)
            # Stop early if all groups are full.
            if all(len(selected[g]) >= n_per_class for g in range(n_groups)):
                break

    # Gather all indices and shuffle them.
    result = []
    for g in range(n_groups):
        result.extend(selected[g])
    rng.shuffle(result)
    print("Ending sampling")

    return np.array(result, dtype=int)


def load_normals_from_h5(h5_path):
    """
    Load normals stored in an HDF5 file of the form:
      points/
        coordinates : (N,3) float32
        normals     : (N,3) float32

    Parameters:
      h5_path (str): path to the .h5 file

    Returns:
      normals (np.ndarray): array of shape (N,3) with dtype float32
    """
    with h5py.File(h5_path, "r") as hf:
        # Move to the "points" group.
        pts_grp = hf["points"]
        # Read the full "normals" dataset directly.
        normals = pts_grp["normals"][:]  # shape (N,3), dtype float32
    return normals





def merge_h5_files_fast(input_paths, output_path):
    """
    Quickly merge several HDF5 files containing:
      - a "points" group with point subgroups
      - an optional "global_stats" group containing the four attributes:
          dist_min, dist_max, angle_min, angle_max

    This function:
      1. Copies every "points" subgroup from each source file
         into output_path/"points" under a new unique index.
      2. Aggregates the global_stats from each source file to produce
         the merged bounds {dist_min, dist_max, angle_min, angle_max},
         then stores them in output_path/"global_stats".

    Note: copying the "points" subgroups is done at the C/HDF5 level
    via h5py.Group.copy(), without loading the data into RAM.

    Parameters:
      input_paths (Iterable[str]): paths to the HDF5 files to merge.
      output_path (str): path of the output HDF5 file.
    """
    # Aggregated global statistics initialized to extreme values.
    agg_stats = {
       'dist_min': float('inf'),
        'dist_max': float('-inf'),
        'angle_min': float('inf'),
        'angle_max': float('-inf'),
        
        'blur_min': float('inf'),
        'blur_max': float('-inf'),
        'contra_min': float('inf'),
        'contra_max': float('-inf'),
        
        'snr_min': float('inf'),
        'snr_max': float('-inf'),
        'sat_min': float('inf'),
        'sat_max': float('-inf')
    }

    # Step 1: copy all "points/*" subgroups.
    with h5py.File(output_path, "w") as hf_out:
        pts_out = hf_out.create_group("points")
        next_idx = 0

        for src_path in input_paths:
            with h5py.File(src_path, "r") as hf_in:
                # 1.a. Copy the subgroups under "points", if present.
                if "points" in hf_in:
                    pts_in = hf_in["points"]
                    for child_name in pts_in:
                        pts_in.copy(child_name, pts_out, name=str(next_idx))
                        next_idx += 1

                # 1.b. Read and aggregate the "global_stats" group if present.

                for key in ("dist_min", "dist_max", "angle_min", "angle_max", 'blur_min','blur_max','contra_min','contra_max','snr_min','snr_max','sat_min','sat_max'):
                            val = hf_in.attrs[key][()]
                            if key.endswith("_min"):
                                agg_stats[key] = min(agg_stats[key], float(val))
                                hf_out.attrs[key] = agg_stats[key]
                            else:  # "_max" key
                                agg_stats[key] = max(agg_stats[key], float(val))
                                hf_out.attrs[key] = agg_stats[key]


 

def  count_points_in_h5(h5_path):
    """
    Open the HDF5 file and return the number of stored points.
    Each point is assumed to be represented by one subgroup under "points/".
    """
    with h5py.File(h5_path, "r") as hf:
        if "points" not in hf:
            return 0
        return len(hf["points"].keys())



def duplictae_h5(original_file,duplicate_file):


    # Open the original HDF5 file for reading
    with h5py.File(original_file, 'r') as hf:
        
        # Create a new HDF5 file to duplicate
        with h5py.File(duplicate_file, 'w') as hf_duplicate:
            
            # Iterate through the groups in 'points'
            for i in hf['points']:
                grp = hf['points'][i]
                
                # Create a new group for each point in the duplicate file
                grp_duplicate = hf_duplicate.create_group(f"points/{i}")
                
                # Copy the 'coordinates' dataset as is
                coords = grp['coordinates'][:]
                grp_duplicate.create_dataset("coordinates", data=coords, dtype="float32")
                
                # Copy the 'pixel_coords' dataset as is
                pixels = grp['pixel_coords'][:10]
                grp_duplicate.create_dataset("pixel_coords", data=pixels, dtype="int32")
                
                # Copy 'ground_truth' dataset as is
                gt = grp['ground_truth'][()]
                grp_duplicate.create_dataset("ground_truth", data=gt, dtype="long")
                
                # Copy 'visibility', but keep only the first 10 elements
                visibility = grp['visibility'][:10]  # Only the first 10 elements
                grp_duplicate.create_dataset("visibility", data=visibility, dtype="float32")
                
                # Copy 'image_ids', but keep only the first 10 elements
                image_ids = grp['image_ids'][:10]  # Only the first 10 elements
                grp_duplicate.create_dataset("image_ids", data=image_ids)
                
                # Copy 'logit_vectors', but keep only the first 10 elements
                logit_vectors = grp['logit_vectors'][:10]  # Only the first 10 elements
                grp_duplicate.create_dataset("logit_vectors", data=logit_vectors, dtype="float32")
    
    print(f"file {original_file} duplicate successfully.")