visualize outputs

Author: pattonw

examples/cremi/cremi.py

@@ -113,204 +113,63 @@
 # Let's visualize the results
 
 # %%
-
 import matplotlib.pyplot as plt
-
-fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-ax[0].imshow(predict_cremi.in_data.array("r")[100], cmap="gray")
-ax[0].set_title("Raw")
-ax[1].imshow(predict_cremi.out_data[0].array("r")[:3, 100].transpose(1, 2, 0))
-ax[1].set_title("LSDs")
-ax[2].imshow(predict_cremi.out_data[1].array("r")[3:6, 100].transpose(1, 2, 0))
-ax[2].set_title("Affinities")
-plt.show()
-
-# %% [markdown]
-# Now we can convert the results to a segmentation. We will run mutex watershed on the affinities in a multi step process.
-# 1) Local fragment extraction - This step runs blockwise and generates fragments from the affinities. For each fragment we save a node in a graph with attributes such as its spatial position and size.
-# 2) Edge extraction - This step runs blockwise and computes mean affinities between fragments, adding edges to the fragment graph.
-# 3) Graph Mutex Watershed - This step runs on the fragment graph, and creates a lookup table from fragment id -> segment id.
-# 4) Relabel fragments - This step runs blockwise and creates the final segmentation.
-
-# %%
-from volara.blockwise import AffAgglom, ExtractFrags, GraphMWS, Relabel
-from volara.datasets import Labels
-from volara.dbs import SQLite
-from volara.lut import LUT
-
-# %% [markdown]
-# First lets define the graph and arrays we are going to use.
-
-# because our graph is in an sql database, we need to define a schema with column names and types
-# for node and edge attributes.
-# For nodes: The defaults such as "id", "position", and "size" are already defined
-# so we only need to define the additional attributes, in this case we have no additional node attributes.
-# For edges: The defaults such as "id", "u", "v" are already defined, so we are only adding the additional
-# attributes "xy_aff", "z_aff", and "lr_aff" for saving the mean affinities between fragments.
-
-# %%
-fragments_graph = SQLite(
-    path="sample_A+_20160601.zarr/fragments.db",
-    edge_attrs={"xy_aff": "float", "z_aff": "float", "lr_aff": "float"},
-)
-fragments_dataset = Labels(store="sample_A+_20160601.zarr/fragments")
-segments_dataset = Labels(store="sample_A+_20160601.zarr/segments")
-
-# %% [markdown]
-# Now we define the tasks with the parameters we want to use.
-
-# %%
-
-# Generate fragments in blocks
-extract_frags = ExtractFrags(
-    db=fragments_graph,
-    affs_data=affs_dataset,
-    frags_data=fragments_dataset,
-    block_size=min_output_shape,
-    context=Coordinate(3, 6, 6) * 2,  # A bit larger than the longest affinity
-    bias=[-0.6] + [-0.4] * 2 + [-0.6] * 2 + [-0.8] * 2,
-    strides=(
-        [Coordinate(1, 1, 1)] * 3
-        + [Coordinate(1, 3, 3)] * 2  # We use larger strides for larger affinities
-        + [Coordinate(1, 6, 6)] * 2  # This is to avoid excessive splitting
-    ),
-    randomized_strides=True,  # converts strides to probabilities of sampling affinities (1/prod(stride))
-    remove_debris=64,  # remove excessively small fragments
-    num_workers=4,
-)
-
-# Generate agglomerated edge scores between fragments via mean affinity accross all edges connecting two fragments
-aff_agglom = AffAgglom(
-    db=fragments_graph,
-    affs_data=affs_dataset,
-    frags_data=fragments_dataset,
-    block_size=min_output_shape,
-    context=Coordinate(3, 6, 6) * 1,
-    scores={
-        "z_aff": affs_dataset.neighborhood[0:1],
-        "xy_aff": affs_dataset.neighborhood[1:3],
-        "lr_aff": affs_dataset.neighborhood[3:],
-    },
-    num_workers=4,
-)
-
-# Run mutex watershed again, this time on the fragment graph with agglomerated edges
-# instead of the voxel graph of affinities
-lut = LUT(path="sample_A+_20160601.zarr/lut.npz")
-total_roi = raw_dataset.array("r").roi
-graph_mws = GraphMWS(
-    db=fragments_graph,
-    lut=lut,
-    weights={"xy_aff": (1, -0.4), "z_aff": (1, -0.6), "lr_aff": (1, -0.6)},
-    roi=(total_roi.offset, total_roi.shape),
-)
-
-# Relabel the fragments into segments
-relabel = Relabel(
-    lut=lut,
-    frags_data=fragments_dataset,
-    seg_data=segments_dataset,
-    block_size=min_output_shape,
-    num_workers=4,
-)
-
-pipeline = extract_frags + aff_agglom + graph_mws + relabel
-pipeline.run_blockwise(multiprocessing=True)
-
-# %% [markdown]
-# Let's visualize
-#
-# If you are following through on your own, I highly recommend installing `funlib.show.neuroglancer`, and
-# running the command line tool via `neuroglancer -d sample_A+_20160601.zarr/*` to visualize the results in
-# neuroglancer.
-#
-# For the purposes of visualizing here, we will make a simple gif
-
-
-# %%
 import matplotlib.animation as animation
-import matplotlib.pyplot as plt
-import numpy as np
-from matplotlib.colors import ListedColormap
-
-fragments = fragments_dataset.array("r")[:, ::2, ::2]
-segments = segments_dataset.array("r")[:, ::2, ::2]
-raw = raw_dataset.array("r")[:, ::2, ::2]
-
-# Get unique labels
-unique_labels = set(np.unique(fragments)) | set(np.unique(segments))
-num_labels = len(unique_labels)
-
-
-def random_color(label):
-    rs = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(label)))
-    return np.array((rs.random(), rs.random(), rs.random()))
 
 
-# Generate random colors for each label
-random_fragment_colors = [random_color(label) for label in range(num_labels)]
-
-# Create a colormap
-cmap_labels = ListedColormap(random_fragment_colors)
-
-# Map labels to indices for the colormap
-label_to_index = {label: i for i, label in enumerate(unique_labels)}
-indexed_fragments = np.vectorize(label_to_index.get)(fragments)
-indexed_segments = np.vectorize(label_to_index.get)(segments)
-
-fig, axes = plt.subplots(1, 3, figsize=(18, 8))
+fig, axes = plt.subplots(1, 3, figsize=(14, 8))
 
 ims = []
-for i, (raw_slice, fragments_slice, segments_slice) in enumerate(
-    zip(raw, indexed_fragments, indexed_segments)
+for i, (raw_slice, affs_slice, lsd_slice) in enumerate(
+    zip(
+        raw_dataset.array("r")[:],
+        affs_dataset.array("r")[:].transpose([1, 0, 2, 3]),
+        lsds_dataset.array("r")[:].transpose([1, 0, 2, 3]),
+    )
 ):
     # Show the raw data
     if i == 0:
         im_raw = axes[0].imshow(raw_slice, cmap="gray")
         axes[0].set_title("Raw")
-        im_fragments = axes[1].imshow(
-            fragments_slice,
-            cmap=cmap_labels,
+        im_affs_long = axes[1].imshow(
+            affs_slice[[0, 5, 6]].transpose([1, 2, 0]),
             vmin=0,
-            vmax=num_labels,
+            vmax=255,
             interpolation="none",
         )
-        axes[1].set_title("Fragments")
-        im_segments = axes[2].imshow(
-            segments_slice,
-            cmap=cmap_labels,
+        axes[1].set_title("Affs (0, 5, 6)")
+        im_lsd = axes[2].imshow(
+            lsd_slice[:3].transpose([1, 2, 0]),
             vmin=0,
-            vmax=num_labels,
+            vmax=255,
             interpolation="none",
         )
-        axes[2].set_title("Segments")
+        axes[2].set_title("LSDs (0, 1, 2)")
     else:
-        im_raw = axes[0].imshow(raw_slice, animated=True, cmap="gray")
-        im_fragments = axes[1].imshow(
-            fragments_slice,
-            cmap=cmap_labels,
+        im_raw = axes[0].imshow(raw_slice, cmap="gray", animated=True)
+        axes[0].set_title("Raw")
+        im_affs_long = axes[1].imshow(
+            affs_slice[[0, 5, 6]].transpose([1, 2, 0]),
             vmin=0,
-            vmax=num_labels,
+            vmax=255,
             interpolation="none",
             animated=True,
         )
-        im_segments = axes[2].imshow(
-            segments_slice,
-            cmap=cmap_labels,
+        axes[1].set_title("Affs (0, 5, 6)")
+        im_lsd = axes[2].imshow(
+            lsd_slice[:3].transpose([1, 2, 0]),
             vmin=0,
-            vmax=num_labels,
+            vmax=255,
             interpolation="none",
             animated=True,
         )
-    ims.append([im_raw, im_fragments, im_segments])
+        axes[2].set_title("LSDs (0, 1, 2)")
+    ims.append([im_raw, im_affs_long, im_lsd])
 
 ims = ims + ims[::-1]
 ani = animation.ArtistAnimation(fig, ims, blit=True)
-ani.save("_static/cremi/segmentation.gif", writer="pillow", fps=10)
+ani.save("_static/cremi/outputs.gif", writer="pillow", fps=10)
 plt.close()
 
 # %% [markdown]
-# The final segmentation is shown below. Obviously this is not a great segmentation, but it is
-# reasonably good for a model small enough to process a CREMI dataset in 20 minutes on a github
-# action.
-# ![segmentation](_static/cremi/segmentation.gif)
+# ![segmentation](_static/cremi/outputs.gif)