fix predicting new channels in an existing store (#57)

mehta-lab · Sep 29, 2023 · 2464a95 · 2464a95
1 parent a57e8c1
commit 2464a95
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/viscy/light/predict_writer.py b/viscy/light/predict_writer.py
@@ -78,10 +78,11 @@ def on_predict_start(self, trainer: Trainer, pl_module: LightningModule) -> None
                     "Cannot write input to an existing store. Aborting."
                 )
             else:
+                with open_ome_zarr(self.output_store, mode="r+") as plate:
+                    for _, pos in plate.positions():
+                        for ch in prediction_channel:
+                            pos.append_channel(ch, resize_arrays=True)
                 self.plate = open_ome_zarr(self.output_store, mode="r+")
-                for _, pos in self.plate.positions():
-                    for ch in prediction_channel:
-                        pos.append_channel(ch, resize_arrays=True)
         else:
             channel_names = prediction_channel
             if self.write_input: