Skip to content

Commit

Permalink
added support for reading and loading stereo events
Browse files Browse the repository at this point in the history
  • Loading branch information
@TjarkMiener
TjarkMiener committed Dec 3, 2024
1 parent 76bf4d4 commit 1f8f9cc
Show file tree
Hide file tree
Showing 2 changed files with 222 additions and 117 deletions.
159 changes: 100 additions & 59 deletions dl1_data_handler/loader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,32 +60,20 @@ def __init__(
self.stack_telescope_images = stack_telescope_images
self.sort_by_intensity = sort_by_intensity

# Get the input shape for the convolutional neural network
if self.DLDataReader.__class__.__name__ == "DLImageReader":
self.channel_shape = len(self.DLDataReader.channels)
elif self.DLDataReader.__class__.__name__ == "DLWaveformReader":
self.channel_shape = self.DLDataReader.sequence_length


# Set the input shape based on the mode of the DLDataReader
if self.DLDataReader.mode == "mono":
self.image_shape = self.DLDataReader.image_mappers[self.DLDataReader.cam_name].image_shape
self.input_shape = (
len(self.DLDataReader.tel_ids),
self.image_shape,
self.image_shape,
self.channel_shape,
)
self.input_shape = self.DLDataReader.input_shape
elif self.DLDataReader.mode == "stereo":
self.input_shape = self.DLDataReader.input_shape[
list(self.DLDataReader.selected_telescopes)[0]
]
# Reshape inputs into proper dimensions for the stereo analysis with stacked images
if self.stack_telescope_images:
# Reshape inputs into proper dimensions for the stereo analysis with stacked images
self.input_shape = (
self.image_shape,
self.image_shape,
len(self.DLDataReader.tel_ids) * self.channel_shape,
self.input_shape[1],
self.input_shape[2],
self.input_shape[0] * self.input_shape[3],
)
else:
self.input_shape = (110, 110, 2)


def __len__(self):
"""
Expand Down Expand Up @@ -128,50 +116,103 @@ def __getitem__(self, index):
A tuple containing the input data as features and the corresponding labels.
"""
# Generate indices of the batch
batch_indices = self.indices[index * self.batch_size : (index + 1) * self.batch_size]
batch_indices = self.indices[
index * self.batch_size : (index + 1) * self.batch_size
]
labels = {}
if self.DLDataReader.mode == "mono":
batch = self.DLDataReader.mono_batch_generation(
batch_indices=batch_indices,
)
batch = self.DLDataReader.generate_mono_batch(batch_indices)
# Retrieve the telescope images and store in the features dictionary
features = {"input": batch["features"].data}
if "type" in self.tasks:
labels["type"] = to_categorical(
batch["true_shower_primary_class"].data,
num_classes=2,
)
# Temp fix till keras support class weights for multiple outputs or I wrote custom loss
# https://github.com/keras-team/keras/issues/11735
if len(self.tasks) == 1:
labels = to_categorical(
batch["true_shower_primary_class"].data,
num_classes=2,
)
if "energy" in self.tasks:
labels["energy"] = batch["log_true_energy"].data
if "direction" in self.tasks:
labels["direction"] = np.stack(
(
batch["spherical_offset_az"].data,
batch["spherical_offset_alt"].data,
batch["angular_separation"].data,
),
axis=1,
)
elif self.DLDataReader.mode == "stereo":
batch = self.DLDataReader.stereo_batch_generation(
batch_indices=batch_indices,
)
batch = self.DLDataReader.generate_stereo_batch(batch_indices)
batch_grouped = batch.group_by(["obs_id", "event_id", "tel_type_id"])
# Sort events based on their telescope types by the hillas intensity in a given batch if requested
if self.sort_by_intensity:
for batch_grouped in batch_grouped.groups:
batch_grouped.sort(["hillas_intensity"], reverse=True)
print(batch_grouped)

if self.stack_telescope_images:
print(features)
# Generate the labels for each task
labels = {}
if "type" in self.tasks:
labels["type"] = to_categorical(
batch["true_shower_primary_class"].data,
num_classes=2,
)
# Temp fix till keras support class weights for multiple outputs or I wrote custom loss
# https://github.com/keras-team/keras/issues/11735
if len(self.tasks) == 1:
labels = to_categorical(
batch["true_shower_primary_class"].data,
features = []
true_shower_primary_class = []
log_true_energy = []
spherical_offset_az, spherical_offset_alt, angular_separation = [], [], []
for group_element in batch_grouped.groups:
if self.sort_by_intensity:
# Sort images by the hillas intensity in a given batch if requested
group_element.sort(["hillas_intensity"], reverse=True)
# Stack the telescope images for stereo analysis
if self.stack_telescope_images:
# Retrieve the telescope images
plain_features = group_element["features"].data
# Stack the telescope images along the last axis
stacked_features = np.concatenate(
[
plain_features[i]
for i in range(plain_features.shape[0])
],
axis=-1,
)
# Append the stacked images to the features list
# shape: (batch_size, image_shape, image_shape, n_channels * n_tel)
features.append(stacked_features)
else:
# Append the plain images to the features list
# shape: (batch_size, n_tel, image_shape, image_shape, n_channels)
features.append(group_element["features"].data)
# Retrieve the labels for the tasks
# FIXME: This won't work for divergent pointing directions
if "type" in self.tasks:
true_shower_primary_class.append(group_element["true_shower_primary_class"].data[0])
if "energy" in self.tasks:
log_true_energy.append(group_element["log_true_energy"].data[0])
if "direction" in self.tasks:
spherical_offset_az.append(group_element["spherical_offset_az"].data[0])
spherical_offset_alt.append(group_element["spherical_offset_alt"].data[0])
angular_separation.append(group_element["angular_separation"].data[0])
# Store the labels in the labels dictionary
if "type" in self.tasks:
labels["type"] = to_categorical(
np.array(true_shower_primary_class),
num_classes=2,
)
if "energy" in self.tasks:
labels["energy"] = batch["log_true_energy"].data
if "direction" in self.tasks:
labels["direction"] = np.stack(
(
batch["spherical_offset_az"].data,
batch["spherical_offset_alt"].data,
batch["angular_separation"].data,
),
axis=1,
)
# Temp fix till keras support class weights for multiple outputs or I wrote custom loss
# https://github.com/keras-team/keras/issues/11735
if len(self.tasks) == 1:
labels = to_categorical(
np.array(true_shower_primary_class),
num_classes=2,
)
if "energy" in self.tasks:
labels["energy"] = np.array(log_true_energy)
if "direction" in self.tasks:
labels["direction"] = np.stack(
(
np.array(spherical_offset_az),
np.array(spherical_offset_alt),
np.array(angular_separation),
),
axis=1,
)
# Store the fatures in the features dictionary
features = {"input": np.array(features)}
# Temp fix for supporting keras2 & keras3
if int(keras.__version__.split(".")[0]) >= 3:
features = features["input"]
Expand Down
Loading

0 comments on commit 1f8f9cc

Please sign in to comment.