LivingOptics/hyperpspectral-forensics

import os
import numpy as np
import matplotlib.pyplot as plt
from lo_dataset_reader import DatasetReader, spectral_coordinate_indices_in_mask, rle_to_mask

os.environ["QT_QPA_PLATFORM"] = "xcb"

dataset_path = "/path/to/dataset"
dataset = DatasetReader(dataset_path, display_fig=True)

for idx, ((info, scene, spectra, unit, images_extern), (converted_spectra, converted_unit), annotations, library_spectra, labels) in enumerate(dataset):
    for ann_idx, annotation in enumerate(annotations):
        annotation["labels"] = labels

        # Visualise the annotation on the scene
        dataset.save_annotation_visualisation(scene, annotation, images_extern, ann_idx)

        # Get spectrum stats from annotation
        stats = annotation.get("extern", {}).get("stats", {})
        label = stats.get("category")
        mean_radiance_spectrum = stats.get("mean_radiance_spectrum")
        mean_reflectance_spectrum = stats.get("mean_reflectance_spectrum")

        # Get mask and spectral indices
        mask = rle_to_mask(annotation["segmentation"], scene.shape)
        spectral_indices = spectral_coordinate_indices_in_mask(mask, info.sampling_coordinates)

        # Extract spectra and converted spectra
        spec = spectra[spectral_indices, :]
        if converted_spectra is not None:
            conv_spec = converted_spectra[spectral_indices, :]
        else:
            conv_spec = None

        # X-axis based on band index or wavelengths (optional)
        x = np.arange(spec.shape[1])
        if stats.get("wavelength_min") is not None and stats.get("wavelength_max") is not None:
            x = np.linspace(stats["wavelength_min"], stats["wavelength_max"], spec.shape[1])

        # Determine plot layout
        if converted_spectra is not None:
            fig, axs = plt.subplots(2, 2, figsize=(12, 8))
            axs_top = axs[0]
            axs_bottom = axs[1]
        else:
            fig, axs_top = plt.subplots(1, 2, figsize=(12, 4))
            print(f"Warning: No converted_spectra for annotation '{label}'")

        unit_label = unit.capitalize() if unit else "Radiance"

        # (1,1) Individual spectra
        for s in spec:
            axs_top[0].plot(x, s, alpha=0.3)
        axs_top[0].set_title(f"{unit_label.capitalize()} Spectra")
        axs_top[0].set_xlabel("Wavelength")
        axs_top[0].set_ylabel(f"{unit_label.capitalize()}")

        # (1,2) Mean + Min/Max (Before conversion)
        if mean_radiance_spectrum is not None:
            spec_min = np.min(spec, axis=0)
            spec_max = np.max(spec, axis=0)
            axs_top[1].fill_between(x, spec_min, spec_max, color='lightblue', alpha=0.5, label='Min-Max Range')
            axs_top[1].plot(x, mean_radiance_spectrum, color='blue', label=f'Mean {unit_label.capitalize()}')
            axs_top[1].set_title(f"Extern Mean ± Range ({unit_label.capitalize()})")
            axs_top[1].set_xlabel("Wavelength")
            axs_top[1].set_ylabel(f"{unit_label.capitalize()}")
            axs_top[1].legend()

        # (2,1) and (2,2) Only if converted_spectra is available
        if converted_spectra is not None and conv_spec is not None:
            for s in conv_spec:
                axs_bottom[0].plot(x, s, alpha=0.3)
            axs_bottom[0].set_title(f"{converted_unit} Spectra")
            axs_bottom[0].set_xlabel("Wavelength")
            axs_bottom[0].set_ylabel(f"{converted_unit}")

            if mean_reflectance_spectrum is not None:
                conv_min = np.min(conv_spec, axis=0)
                conv_max = np.max(conv_spec, axis=0)
                axs_bottom[1].fill_between(x, conv_min, conv_max, color='lightgreen', alpha=0.5, label='Min-Max Range')
                axs_bottom[1].plot(x, mean_reflectance_spectrum, color='green', label=f'Mean {converted_unit}')
                axs_bottom[1].set_title(f"Extern Mean ± Range ({converted_unit})")
                axs_bottom[1].set_xlabel("Wavelength")
                axs_bottom[1].set_ylabel(f"{converted_unit}")
                axs_bottom[1].legend()

        fig.suptitle(f"Annotation {label}", fontsize=16)
        plt.tight_layout()
        plt.show()