Merge pull request #8 from AllenNeuralDynamics/run-without-cellpose

Run without cellpose

code/extraction.py

@@ -780,12 +780,11 @@ def contour_video(
             raw_r = []
         keys = []
 
-    cellpose_path = str(next(Path(args.tmp_dir).rglob("cellpose.npz")))
+    cellpose_path = str(next(Path(args.tmp_dir).rglob("cellpose.npz"), ""))
     ops_path = str(next(Path(args.tmp_dir).rglob("ops.npy")))
     # write output files
     with (
-        h5py.File(output_dir / "extraction.h5", "w") as f,
-        np.load(cellpose_path) as cp,
+        h5py.File(output_dir / "extraction.h5", "w") as f
     ):
         # traces
         f.create_dataset("traces/corrected", data=traces_corrected, compression="gzip")
@@ -813,8 +812,13 @@ def contour_video(
             "rois/neuropil_coords", data=neuropil_coords, compression="gzip"
         )
         # cellpose
-        for k in cp.keys():
-            f.create_dataset(f"cellpose/{k}", data=cp[k], compression="gzip")
+        if cellpose_path:
+            with np.load(cellpose_path) as cp:
+                for k in cp.keys():
+                    f.create_dataset(f"cellpose/{k}", data=cp[k], compression="gzip")
+        else:
+            logging.warning("No cellpose output found.")
+
         # classifier
         f.create_dataset(f"iscell", data=iscell, dtype="f4")
         # summary images