I2K 2020 webKnossos Scripts

density.py

import wknml
import numpy as np


def calculate_density(trees, bbox, scale):
    assert len(trees) > 1
    bbox_lo = np.array(bbox[0:3])
    bbox_hi = bbox_lo + np.array(bbox[3:6])

    result = 0
    for tree in trees:
        assert len(tree.nodes) > 0
        node = tree.nodes[0]
        if np.all(np.array(node.position) >= bbox_lo) and np.all(
            np.array(node.position) <= bbox_hi
        ):
            result += 1

    return result / np.prod((bbox_hi - bbox_lo) * np.array(scale))


nml = wknml.parse_nml(open("user-bbox.nml", "r"))
print(
    calculate_density(nml.trees, (663, 962, 70, 2100, 1500, 2600), nml.parameters.scale)
    * 1e9,
    "per µm^3",
)

path_length.py

import wknml
import numpy as np


def calculate_path_length_for_tree(tree, scale):
    assert len(tree.nodes) > 1
    result = 0
    for edge in tree.edges:
        a = next(node for node in tree.nodes if node.id == edge.source)
        b = next(node for node in tree.nodes if node.id == edge.target)
        diff_vector = np.array(a.position) - np.array(b.position)
        scaled_diff_vector = diff_vector * scale
        edge_length = np.sqrt(scaled_diff_vector.dot(scaled_diff_vector))
        result += edge_length
    return result


nml = wknml.parse_nml(open("single-tree.nml", "r"))
print(calculate_path_length_for_tree(nml.trees[0], nml.parameters.scale) / 1e3, "µm")

single-tree.nml

<things>
  <meta name="writer" content="nml_helpers.js" />
  <meta name="writerGitCommit" content="3e07a3c1d77dfbf7576c99660b96b2d734dd463a" />
  <meta name="timestamp" content="1606223862730" />
  <meta name="annotationId" content="5fbd002c01000018006f987f" />
  <meta name="username" content="Norman Rzepka" />
  <parameters>
    <experiment name="2012-09-28_ex145_07x2_ROI2017" description="" organization="MPI_Brain_Research" />
    <scale x="11.239999771118164" y="11.239999771118164" z="28" />
    <offset x="0" y="0" z="0" />
    <time ms="1606221866173" />
    <editPosition x="2994" y="4413" z="1571" />
    <editRotation xRot="0" yRot="0" zRot="0" />
    <zoomLevel zoom="0.45" />
  </parameters>
  <thing id="1" color.r="0.6784313725490196" color.g="0.1411764705882353" color.b="0.050980392156862744" color.a="1" name="explorative_2020-11-24_Norman_Rzepka_001" groupId="">
    <nodes>
      <node id="1" radius="112.39999771118164" x="2671" y="4335" z="1727" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223834269" />
      <node id="2" radius="112.39999771118164" x="2752" y="4368" z="1681" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223836432" />
      <node id="3" radius="112.39999771118164" x="2846" y="4382" z="1656" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223839133" />
      <node id="4" radius="112.39999771118164" x="2894" y="4387" z="1634" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223840809" />
      <node id="5" radius="112.39999771118164" x="2915" y="4396" z="1614" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223842239" />
      <node id="6" radius="112.39999771118164" x="2936" y="4401" z="1597" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223843458" />
      <node id="7" radius="112.39999771118164" x="2966" y="4412" z="1586" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223844812" />
      <node id="8" radius="112.39999771118164" x="2994" y="4412" z="1571" rotX="0" rotY="0" rotZ="0" inVp="0" inMag="0" bitDepth="8" interpolation="true" time="1606223846560" />
    </nodes>
    <edges>
      <edge source="1" target="2" />
      <edge source="2" target="3" />
      <edge source="3" target="4" />
      <edge source="4" target="5" />
      <edge source="5" target="6" />
      <edge source="6" target="7" />
      <edge source="7" target="8" />
    </edges>
  </thing>
  <branchpoints>
  </branchpoints>
  <comments>
  </comments>
  <groups>
  </groups>
</things>

train.py

import numpy as np
from tempfile import mkdtemp
from os import path, makedirs
import zipfile
import shutil
import wkw
import wknml
from wkcuber.metadata import write_webknossos_metadata

SHAPE = (256, 256, 30)
SCALE = (5,5,20)


def read_volume_annotation(filename, shape):
  """Reads a webKnossos volume annotation"""
  tmp_dir = mkdtemp()
  with zipfile.ZipFile(filename, 'r') as zip_ref:
    zip_ref.extractall(tmp_dir)
  with zipfile.ZipFile(path.join(tmp_dir, 'data.zip'), 'r') as zip_ref:
    zip_ref.extractall(path.join(tmp_dir, 'segmentation'))
  ds = wkw.Dataset.open(path.join(tmp_dir, 'segmentation', '1'))
  return ds.read((0, 0, 0), shape)[0]


def write_layer_buffered(path_name, data):
  ds = wkw.Dataset.create(path_name, wkw.Header(voxel_type=data.dtype, block_type=wkw.Header.BLOCK_TYPE_LZ4HC))
  buffer = np.zeros((1024, 1024, 1024), dtype=data.dtype)
  buffer[0:SHAPE[0], 0:SHAPE[1], 0:SHAPE[2]] = data
  ds.write((0, 0, 0), buffer)


def write_dataset(filename, scale, raw_data, prediction_data, segmentation_threshold=0.5):
  """Creates a zip file with a dataset that contains the raw, prediction and segmentation data"""
  tmp_dir = mkdtemp()

  makedirs(path.join(tmp_dir, 'color', '1'))
  write_layer_buffered(path.join(tmp_dir, 'color', '1'), raw_data.T)

  makedirs(path.join(tmp_dir, 'prediction', '1'))
  write_layer_buffered(path.join(tmp_dir, 'prediction', '1'), prediction_data.T)

  makedirs(path.join(tmp_dir, 'segmentation', '1'))
  write_layer_buffered(path.join(tmp_dir, 'segmentation', '1'), (prediction_data.T >= segmentation_threshold).astype(np.uint8))

  write_webknossos_metadata(tmp_dir, filename, scale, compute_max_id=True, exact_bounding_box={"topLeft": (0,0,0), "width": SHAPE[0], "height":SHAPE[1], "depth": SHAPE[2]})

  shutil.make_archive(filename, 'zip', path.join(tmp_dir))


def get_model():
  """Creates a dummy model that returns a pre-predicted volume."""
  class Model:
    def fit(self, X, y, **kwargs):
      pass
    def predict(self, X):
      return wkw.Dataset.open(path.join("multitask_segmentation_skull", "prediction", "1")).read((0,0,0), SHAPE).T.astype(np.float32) / 255.0
  return Model()


# Load ground truth
train_data = wkw.Dataset.open(path.join("mri-brain-02", "color", "1")).read((0,0,0), SHAPE).T
train_labels = read_volume_annotation("mri-brain-02__explorational.zip", SHAPE).T

# Loads another dataset for evaluation
eval_data = wkw.Dataset.open(path.join("mri-brain-02", "color", "1")).read((0,0,0), SHAPE).T

# Train a model
model = get_model()
model.fit(train_data, train_labels, batch_size=32, epochs=5)

# Create a prediction
predict_data = model.predict(eval_data)
# Write out the data for upload to webKnossos
write_dataset("mri-brain-02-prediction", SCALE, eval_data, predict_data)

user-bbox.nml

<things>
  <meta name="writer" content="NmlWriter.scala" />
  <meta name="writerGitCommit" content="7f297164fffcfee52e12e31e91e5a7af4a6ea139" />
  <meta name="timestamp" content="1606756336345" />
  <meta name="annotationId" content="5fc526d50100007b00cfac32" />
  <meta name="username" content="Norman Rzepka" />
  <parameters>
    <experiment name="xnh_esrf_mouseCortex_30nm" organization="Demo_Organization" description="" />
    <scale x="30.0" y="30.0" z="30.0" />
    <offset x="0" y="0" z="0" />
    <time ms="1606756053583" />
    <editPosition x="1689" y="1426" z="1927" />
    <editRotation xRot="0.0" yRot="0.0" zRot="0.0" />
    <zoomLevel zoom="8.140274938683982" />
    <activeNode id="33" />
    <userBoundingBox id="0" name="user bounding box 0" isVisible="true" color.r="0.13725490868091583" color.g="0.3843137323856354" color.b="0.843137264251709" color.a="1.0" topLeftX="663" topLeftY="962" topLeftZ="70" width="2100" height="1500" depth="2600" />
  </parameters>
  <thing id="33" color.r="0.8274509906768799" color.g="0.30980393290519714" color.b="0.5176470875740051" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_033">
    <nodes>
      <node id="33" radius="300.0" x="2808" y="1675" z="2380" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756173886" />
    </nodes>
    <edges />
  </thing>
  <thing id="32" color.r="0.8274509906768799" color.g="0.6549019813537598" color.b="0.6549019813537598" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_032">
    <nodes>
      <node id="32" radius="300.0" x="2808" y="634" z="652" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756172637" />
    </nodes>
    <edges />
  </thing>
  <thing id="31" color.r="0.07058823853731155" color.g="0.6549019813537598" color.b="0.3803921639919281" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_031">
    <nodes>
      <node id="31" radius="300.0" x="2725" y="1919" z="1870" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756170258" />
    </nodes>
    <edges />
  </thing>
  <thing id="30" color.r="0.27450981736183167" color.g="1.0" color.b="0.48235294222831726" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_030">
    <nodes>
      <node id="30" radius="300.0" x="2725" y="2263" z="1790" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756169295" />
    </nodes>
    <edges />
  </thing>
  <thing id="29" color.r="1.0" color.g="0.48235294222831726" color.b="0.3803921639919281" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_029">
    <nodes>
      <node id="29" radius="300.0" x="2725" y="2580" z="719" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756168557" />
    </nodes>
    <edges />
  </thing>
  <thing id="28" color.r="0.929411768913269" color.g="1.0" color.b="0.6901960968971252" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_028">
    <nodes>
      <node id="28" radius="300.0" x="2591" y="2556" z="220" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756165363" />
    </nodes>
    <edges />
  </thing>
  <thing id="27" color.r="0.24313725531101227" color.g="0.0" color.b="0.10196078568696976" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_027">
    <nodes>
      <node id="27" radius="300.0" x="2486" y="2267" z="1794" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756161372" />
    </nodes>
    <edges />
  </thing>
  <thing id="26" color.r="0.7254902124404907" color.g="0.30980393290519714" color.b="0.8274509906768799" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_026">
    <nodes>
      <node id="26" radius="300.0" x="2486" y="1915" z="1825" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756160703" />
    </nodes>
    <edges />
  </thing>
  <thing id="25" color.r="0.8274509906768799" color.g="1.0" color.b="0.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_025">
    <nodes>
      <node id="25" radius="300.0" x="2486" y="1905" z="1633" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756160206" />
    </nodes>
    <edges />
  </thing>
  <thing id="24" color.r="0.0" color.g="0.27450981736183167" color.b="0.5843137502670288" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_024">
    <nodes>
      <node id="24" radius="300.0" x="2392" y="1901" z="1671" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756148907" />
    </nodes>
    <edges />
  </thing>
  <thing id="23" color.r="0.0" color.g="0.30980393290519714" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_023">
    <nodes>
      <node id="23" radius="300.0" x="2376" y="1278" z="2247" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756147470" />
    </nodes>
    <edges />
  </thing>
  <thing id="22" color.r="0.6196078658103943" color.g="0.0" color.b="0.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_022">
    <nodes>
      <node id="22" radius="300.0" x="2376" y="1713" z="2097" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756145997" />
    </nodes>
    <edges />
  </thing>
  <thing id="21" color.r="0.4470588266849518" color.g="0.7843137383460999" color.b="0.48235294222831726" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_021">
    <nodes>
      <node id="21" radius="300.0" x="2376" y="2733" z="167" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756144792" />
    </nodes>
    <edges />
  </thing>
  <thing id="20" color.r="0.6196078658103943" color.g="0.7568627595901489" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_020">
    <nodes>
      <node id="20" radius="300.0" x="2376" y="1901" z="59" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756143870" />
    </nodes>
    <edges />
  </thing>
  <thing id="19" color.r="0.4470588266849518" color.g="0.9647058844566345" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_019">
    <nodes>
      <node id="19" radius="300.0" x="2376" y="1549" z="551" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756142895" />
    </nodes>
    <edges />
  </thing>
  <thing id="18" color.r="0.5176470875740051" color.g="0.4470588266849518" color.b="0.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_018">
    <nodes>
      <node id="18" radius="300.0" x="2376" y="1240" z="261" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756141823" />
    </nodes>
    <edges />
  </thing>
  <thing id="17" color.r="0.929411768913269" color.g="0.8470588326454163" color.b="0.6784313917160034" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_017">
    <nodes>
      <node id="17" radius="300.0" x="2344" y="588" z="495" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756139983" />
    </nodes>
    <edges />
  </thing>
  <thing id="16" color.r="1.0" color.g="0.239215686917305" color.b="0.5882353186607361" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_016">
    <nodes>
      <node id="16" radius="300.0" x="2225" y="1281" z="2216" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756137928" />
    </nodes>
    <edges />
  </thing>
  <thing id="15" color.r="0.0" color.g="0.4000000059604645" color.b="0.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_015">
    <nodes>
      <node id="15" radius="300.0" x="2106" y="2937" z="2275" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756134803" />
    </nodes>
    <edges />
  </thing>
  <thing id="14" color.r="0.5490196347236633" color.g="0.2705882489681244" color.b="0.07058823853731155" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_014">
    <nodes>
      <node id="14" radius="300.0" x="2106" y="1889" z="73" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756133765" />
    </nodes>
    <edges />
  </thing>
  <thing id="13" color.r="1.0" color.g="1.0" color.b="0.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_013">
    <nodes>
      <node id="13" radius="300.0" x="2106" y="472" z="1200" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756132946" />
    </nodes>
    <edges />
  </thing>
  <thing id="12" color.r="0.239215686917305" color.g="0.7803921699523926" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_012">
    <nodes>
      <node id="12" radius="300.0" x="2106" y="580" z="1748" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756132022" />
    </nodes>
    <edges />
  </thing>
  <thing id="11" color.r="0.45098039507865906" color.g="0.239215686917305" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_011">
    <nodes>
      <node id="11" radius="300.0" x="2090" y="1677" z="2150" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756130752" />
    </nodes>
    <edges />
  </thing>
  <thing id="10" color.r="0.4117647111415863" color.g="0.8313725590705872" color.b="0.0313725508749485" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_010">
    <nodes>
      <node id="10" radius="300.0" x="1978" y="559" z="45" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756126212" />
    </nodes>
    <edges />
  </thing>
  <thing id="9" color.r="0.7215686440467834" color.g="0.05882352963089943" color.b="0.8196078538894653" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_009">
    <nodes>
      <node id="9" radius="300.0" x="1752" y="674" z="2380" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756119333" />
    </nodes>
    <edges />
  </thing>
  <thing id="8" color.r="0.9647058844566345" color.g="0.6196078658103943" color.b="0.8627451062202454" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_008">
    <nodes>
      <node id="8" radius="300.0" x="1752" y="1743" z="2153" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756118388" />
    </nodes>
    <edges />
  </thing>
  <thing id="7" color.r="0.9647058844566345" color.g="0.07058823853731155" color.b="0.3803921639919281" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_007">
    <nodes>
      <node id="7" radius="300.0" x="1752" y="2464" z="1954" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756117498" />
    </nodes>
    <edges />
  </thing>
  <thing id="6" color.r="0.0784313753247261" color.g="0.4000000059604645" color.b="0.4901960790157318" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_006">
    <nodes>
      <node id="6" radius="300.0" x="1752" y="1827" z="213" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756116470" />
    </nodes>
    <edges />
  </thing>
  <thing id="5" color.r="1.0" color.g="0.8392156958580017" color.b="0.239215686917305" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_005">
    <nodes>
      <node id="5" radius="300.0" x="1752" y="1628" z="1654" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756115620" />
    </nodes>
    <edges />
  </thing>
  <thing id="4" color.r="0.0" color.g="1.0" color.b="0.7568627595901489" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_004">
    <nodes>
      <node id="4" radius="300.0" x="1752" y="1604" z="1392" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756114599" />
    </nodes>
    <edges />
  </thing>
  <thing id="3" color.r="0.8509804010391235" color.g="0.5215686559677124" color.b="0.07058823853731155" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_003">
    <nodes>
      <node id="3" radius="300.0" x="1752" y="103" z="2289" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756113548" />
    </nodes>
    <edges />
  </thing>
  <thing id="2" color.r="0.0" color.g="0.0" color.b="1.0" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_002">
    <nodes>
      <node id="2" radius="300.0" x="1752" y="65" z="1553" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756112665" />
    </nodes>
    <edges />
  </thing>
  <thing id="1" color.r="0.6784313917160034" color.g="0.1411764770746231" color.b="0.05098039284348488" color.a="1.0" name="explorative_2020-11-30_Norman_Rzepka_001">
    <nodes>
      <node id="1" radius="300.0" x="1752" y="235" z="1172" rotX="0.0" rotY="270.0" rotZ="0.0" inVp="1" inMag="4" bitDepth="8" interpolation="true" time="1606756111250" />
    </nodes>
    <edges />
  </thing>
  <branchpoints />
  <comments />
  <groups />
  <volume id="1" location="data.zip" />
</things>