Prediction failed using a model trained from scratch

[thao-do]

Describe the bug

Prediction failed using a model trained from scratch

To Reproduce

Steps to reproduce the behavior:

1. Train a model from scratch
2. Go to Prediction & run a prediction on an on-screen image
3. See error

Expected behavior

A segmentation prediction should be generated without any error

Screenshots

N/A

Describe your data (image format, 2D /3D etc.)

Tried for:

• ome-tiff, single-channel, LaminB1 (nucleus)
• ome-tiff, multi-channel, LaminB1 (nucleus)

Environment (please complete the following information):

• Winsegml2 EC2 instance
• OS: Window
• Plugin Version: iterative training fix 11/23/2024
• PyTorch version: -
• GPU? Yes
• CUDA version -

Additional context

Error code below:

Initializing PersistentDataset for train at C:\Users\Administrator\Desktop\SegmenterML-testData\output-test-10\cache\train
Initializing PersistentDataset for val at C:\Users\Administrator\Desktop\SegmenterML-testData\output-test-10\cache\val
INFO: Using 16bit Automatic Mixed Precision (AMP)
INFO: Trainer already configured with model summary callbacks: [<class 'lightning.pytorch.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.
INFO: GPU available: True (cuda), used: True
INFO: TPU available: False, using: 0 TPU cores
INFO: HPU available: False, using: 0 HPUs
INFO: Restoring states from the checkpoint path at C:\Users\Administrator\Desktop\segmenter-home\iterative-training_no-weigth-10\checkpoints\epoch_001.ckpt
INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO: Loaded model weights from the checkpoint at C:\Users\Administrator\Desktop\segmenter-home\iterative-training_no-weigth-10\checkpoints\epoch_001.ckpt
Predicting ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1/1 0:00:01 • 0:00:00 0.00it/s
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File C:\ProgramData\miniconda3\envs\iterative-training\lib\site-packages\allencell_ml_segmenter\core\view.py:94, in View._onLongTaskThreadFinished(self=<allencell_ml_segmenter.prediction.view.PredictionView object>)
     92 self.longTaskThread.deleteLater()
     93 self.progressDialog.close()
---> 94 self.showResults()
        self = <allencell_ml_segmenter.prediction.view.PredictionView object at 0x00000219F1004A60>

File C:\ProgramData\miniconda3\envs\iterative-training\lib\site-packages\allencell_ml_segmenter\prediction\view.py:186, in PredictionView.showResults(self=<allencell_ml_segmenter.prediction.view.PredictionView object>)
    178 for data in stem_to_data.values():
    179     raw_np_data: Optional[np.ndarray] = (
    180         self._img_data_extractor.extract_image_data(
    181             data["raw"], channel=channel
    182         ).np_data
    183     )
    184     seg_np_data: Optional[np.ndarray] = (
    185         self._img_data_extractor.extract_image_data(
--> 186             data["seg"], seg=1
        np = <module 'numpy' from 'C:\\ProgramData\\miniconda3\\envs\\iterative-training\\lib\\site-packages\\numpy\\__init__.py'>
        self = <allencell_ml_segmenter.prediction.view.PredictionView object at 0x00000219F1004A60>
        self._img_data_extractor = <allencell_ml_segmenter.core.image_data_extractor.aics_image_data_extractor.AICSImageDataExtractor object at 0x00000219DDB17CA0>
        data = {'raw': WindowsPath('C:/Users/Administrator/Desktop/SegmenterML-testData/LaminB1-Variance-multiChannels-20241024T144528Z-001/image-2.tif')}
    187         ).np_data
    188     )
    189     if raw_np_data is not None:
    190         self._viewer.add_image(
    191             raw_np_data,
    192             f"[raw] {data['raw'].name}",
    193         )

KeyError: 'seg'