Skip to content

Commit c590869

Browse files
Update processing scripts for compartment annotations
1 parent a18a992 commit c590869

File tree

2 files changed

+111
-0
lines changed

2 files changed

+111
-0
lines changed
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import os
2+
from glob import glob
3+
4+
import h5py
5+
import imageio.v3 as imageio
6+
import napari
7+
import numpy as np
8+
9+
from skimage.measure import label
10+
# from skimage.morphology import remove_small_holes
11+
from tqdm import tqdm
12+
13+
14+
def process_labels(labels):
15+
labels = label(labels)
16+
17+
min_size = 75
18+
ids, sizes = np.unique(labels, return_counts=True)
19+
filter_ids = ids[sizes < min_size]
20+
labels[np.isin(labels, filter_ids)] = 0
21+
22+
# labels = remove_small_holes(labels, area_threshold=min_size)
23+
return labels
24+
25+
26+
def postprocess_annotation(im_path, ann_path, output_folder, view=False):
27+
fname = os.path.basename(im_path)
28+
29+
out_path = os.path.join(output_folder, fname.replace(".tif", ".h5"))
30+
if os.path.exists(out_path):
31+
return
32+
33+
labels = imageio.imread(ann_path)
34+
35+
# Skip empty labels.
36+
if labels.max() == 0:
37+
print("Skipping", im_path)
38+
return
39+
40+
image = imageio.imread(im_path)
41+
labels = process_labels(labels)
42+
43+
if view:
44+
v = napari.Viewer()
45+
v.add_image(image)
46+
v.add_labels(labels)
47+
napari.run()
48+
return
49+
50+
with h5py.File(out_path, "a") as f:
51+
f.create_dataset("data", data=image, compression="gzip")
52+
f.create_dataset("labels/compartments", data=labels, compression="gzip")
53+
54+
55+
def postprocess_annotations(view):
56+
images = sorted(glob("output/images/*.tif"))
57+
annotations = sorted(glob("output/annotations/*.tif"))
58+
59+
output_folder = "output/postprocessed_annotations"
60+
os.makedirs(output_folder, exist_ok=True)
61+
for im, ann in tqdm(zip(images, annotations), total=len(images)):
62+
postprocess_annotation(im, ann, output_folder, view=view)
63+
64+
65+
def main():
66+
postprocess_annotations(view=False)
67+
68+
69+
if __name__ == "__main__":
70+
main()
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import os
2+
3+
from elf.io import open_file
4+
from micro_sam.sam_annotator import annotator_3d, image_folder_annotator
5+
6+
7+
def run_volume_annotation(ds, name):
8+
checkpoint_path = "./checkpoints/compartment_model/best.pt"
9+
10+
tomogram_path = f"./output/{ds}/tomograms/{name}.h5"
11+
embedding_path = f"./output/{ds}/embeddings/{name}.zarr"
12+
assert os.path.exists(embedding_path), embedding_path
13+
14+
with open_file(tomogram_path, "r") as f:
15+
tomogram = f["data"][:]
16+
annotator_3d(tomogram, embedding_path=embedding_path, model_type="vit_b", checkpoint_path=checkpoint_path)
17+
18+
19+
def run_image_annotation():
20+
checkpoint_path = "./checkpoints/compartment_model/best.pt"
21+
image_folder_annotator(
22+
input_folder="output/images", output_folder="output/annotations", pattern="*.tif",
23+
embedding_path="output/embeddings", model_type="vit_b",
24+
checkpoint_path=checkpoint_path
25+
)
26+
27+
28+
def main():
29+
run_image_annotation()
30+
31+
# ds = "09_stem750_66k"
32+
# name = "36859_J1_66K_TS_PS_05_rec_2kb1dawbp_crop"
33+
34+
# ds = "cryo"
35+
# name = "vesicles-64K-LAM12"
36+
37+
# run_annotation(ds, name)
38+
39+
40+
if __name__ == "__main__":
41+
main()

0 commit comments

Comments
 (0)