from coseda.logging_utils import log_print
import h5py
import numpy as np
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt
from typing import Tuple
from coseda.nexus.paths import (
DETECTOR_GROUP_PATH,
LEGACY_MASK_PATH,
NEXUS_MASK_PATH,
normalize_dataset_path,
)
def _ensure_mask_hardlink(h5file: h5py.File, target_path: str, link_path: str) -> None:
"""Create/refresh a hard link from link_path to target_path."""
if link_path in h5file:
del h5file[link_path]
h5file[link_path] = h5file[target_path]
[docs]
def generate_mask(image_path: Path, threshold: int = 127) -> (np.ndarray, np.ndarray):
"""
Load an image (BMP, PNG, JPG), convert to grayscale, and create a binary mask:
valid pixels (above `threshold`) -> 1, invalid (at or below `threshold`) -> 0
Returns the grayscale and mask arrays.
"""
img = Image.open(image_path)
gray = img.convert('L')
arr = np.array(gray)
# Threshold: anything above `threshold` is considered valid
mask = (arr > threshold).astype(np.uint8)
return gray, mask
[docs]
def save_mask_hdf5(mask: np.ndarray,
hdf5_path: Path,
dataset: str = 'mask'):
"""
Add (or overwrite) a mask array in an existing HDF5 file under `dataset`.
- If the file doesn't exist, it will be created.
- If the dataset already exists, only that dataset is replaced.
"""
# ensure parent folder exists
hdf5_path.parent.mkdir(parents=True, exist_ok=True)
alias_path = normalize_dataset_path(dataset)
# open in append mode
with h5py.File(hdf5_path, 'a') as f:
# Ensure canonical NeXus location, then create a hard link for legacy access.
if NEXUS_MASK_PATH in f:
del f[NEXUS_MASK_PATH]
detector_group = f.require_group(DETECTOR_GROUP_PATH.lstrip("/"))
detector_group.create_dataset("mask", data=mask, dtype='uint8')
_ensure_mask_hardlink(f, NEXUS_MASK_PATH, LEGACY_MASK_PATH)
if alias_path not in (NEXUS_MASK_PATH, LEGACY_MASK_PATH):
_ensure_mask_hardlink(f, NEXUS_MASK_PATH, alias_path)
log_print(f"Mask written to '{NEXUS_MASK_PATH}' in {hdf5_path} (shape={mask.shape})")
[docs]
def remove_mask_hdf5(hdf5_path: Path, dataset: str = 'mask'):
"""
Remove the specified mask dataset from the HDF5 file.
If the dataset does not exist, no error is raised.
"""
alias_path = normalize_dataset_path(dataset)
with h5py.File(hdf5_path, 'a') as f:
removed_any = False
for path in {alias_path, LEGACY_MASK_PATH, NEXUS_MASK_PATH}:
if path in f:
del f[path]
log_print(f"→ '{path}' removed from {hdf5_path}")
removed_any = True
if not removed_any:
log_print(f"→ No mask dataset found in {hdf5_path}, nothing to remove.")
[docs]
def load_mask_hdf5(hdf5_path: Path) -> np.ndarray:
"""
Load the mask dataset from HDF5 file (NeXus path preferred).
"""
with h5py.File(hdf5_path, 'r') as f:
if NEXUS_MASK_PATH in f:
return f[NEXUS_MASK_PATH][:]
if LEGACY_MASK_PATH in f:
return f[LEGACY_MASK_PATH][:]
raise KeyError("No mask dataset found at '/entry/instrument/detector/mask' or '/mask'.")
[docs]
def generate_full_mask(shape: Tuple[int, int]) -> np.ndarray:
"""
Generate a binary mask with all pixels valid (1) for a given 2D shape.
Parameters:
shape : tuple of (height, width)
The dimensions of the mask to create.
Returns:
mask : np.ndarray of uint8
2D array of ones with the specified shape.
"""
return np.ones(shape, dtype=np.uint8)
[docs]
def generate_full_mask_hdf5(hdf5_path: Path,
image_dataset: str = 'entry/data/images',
mask_dataset: str = 'mask') -> None:
"""
Inspect the given HDF5 file, read the shape of the specified 2D image stack dataset,
generate a full-valid (all ones) mask of the same [height, width], and save it to the file.
Overwrites any existing mask dataset at `mask_dataset`.
Parameters:
hdf5_path : Path
Path to the HDF5 file containing the image stack.
image_dataset : str
Name of the dataset in the HDF5 file that contains the image stack
(expected shape: (n_frames, height, width) or (height, width)).
mask_dataset : str
Name of the dataset under which to store the generated mask.
"""
# Open in read/write mode
alias_path = normalize_dataset_path(mask_dataset)
with h5py.File(hdf5_path, 'a') as f:
if image_dataset not in f:
raise KeyError(f"Image dataset '{image_dataset}' not found in {hdf5_path}")
data_shape = f[image_dataset].shape
# Determine 2D shape
if len(data_shape) == 3:
_, h, w = data_shape
else:
raise ValueError(f"Unsupported dataset shape {data_shape}.")
# Create full mask
mask = np.ones((h, w), dtype=np.uint8)
if NEXUS_MASK_PATH in f:
del f[NEXUS_MASK_PATH]
detector_group = f.require_group(DETECTOR_GROUP_PATH.lstrip("/"))
detector_group.create_dataset("mask", data=mask, dtype='uint8')
_ensure_mask_hardlink(f, NEXUS_MASK_PATH, LEGACY_MASK_PATH)
if alias_path not in (NEXUS_MASK_PATH, LEGACY_MASK_PATH):
_ensure_mask_hardlink(f, NEXUS_MASK_PATH, alias_path)
log_print(f"Full mask saved to '{NEXUS_MASK_PATH}' in {hdf5_path} with shape {(h, w)}")
[docs]
def plot_comparison(gray: np.ndarray, mask: np.ndarray):
"""
Display original grayscale image and mask side by side for debugging.
"""
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].imshow(gray, cmap='gray')
axes[0].set_title('Original Grayscale')
axes[0].axis('off')
axes[1].imshow(mask, cmap='gray')
axes[1].set_title('Mask (1=valid, 0=invalid)')
axes[1].axis('off')
plt.tight_layout()
plt.show()