precomputed

precomputed.R

library(nat)
library(jsonlite)
library(float)

#' Based on specification
#' https://github.com/google/neuroglancer/tree/master/src/neuroglancer/datasource/precomputed

# write mesh to single res precomputed format
#meshes = read.neurons("../registration/dns_objs21_reg/")

#n_mesh = length(meshes)

save_single_res_info <- function(output_info) {
  jsonlite::write_json(
    list("@type" = "neuroglancer_legacy_mesh",
         "spatial_index"= NA,
         "segment_name_map"="segment_names",
         "segment_properties"="segment_properties"),
    output_info,
    auto_unbox = TRUE,
  )
}

save_single_res_mesh <- function(mesh, output_path, name) {
  fc = file(file.path(output_path, name), "wb")
  writeBin(as.integer(dim(mesh$vb)[[2]]), fc)
  writeBin(as.vector(mesh$vb[1:3,]), fc, size = 4L)
  writeBin(as.integer(mesh$it-1), fc, size = 4L)
  close(fc)
}

save_single_res_manifest <- function(output_path, name) {
  jsonlite::write_json(
    list("fragments" = as.character(name)),
    file.path(output_path, paste0(name, ":0")),
    auto_unbox = FALSE,
  )
}

write_meshes_info <- function(meshes, start_idx,
                              output_seg_names, output_seg_prop) {
  volids <- as.character(
    start_idx:(start_idx + length(meshes) - 1)
  )
  if (is.null(names(meshes)))
    map_ids <- as.list(1:length(meshes))
  else
    map_ids <- as.list(names(meshes))
  names(map_ids) <- volids
  # segment properties
  info <- list(
    "@type"=jsonlite::unbox("neuroglancer_segment_properties")
  )
  mesh_properties = list("id" = jsonlite::unbox("label"),
    "type"= jsonlite::unbox("label"),
    "values"= names(map_ids)
  )
  info$inline = list(
    "ids" = volids,
    "properties" =  list(mesh_properties)
  )
  jsonlite::write_json(
    info,
    file.path(output_seg_prop, "info"),
    auto_unbox = FALSE,
  )
  # segment names
  info <- list(
    "@type"="neuroglancer_segment_name_map",
    "map" = map_ids
  )
  jsonlite::write_json(
    info,
    file.path(output_seg_names, "info"),
    auto_unbox = TRUE,
  )
}

save_single_res_meshes <- function(mesh_list, output_path, start_index = 1) {
  #' TODO check if mesh_list is a nat::neuronlist
  output_meshes = file.path(output_path, "mesh")
  output_info = file.path(output_meshes, "info")
  output_seg_names = file.path(output_meshes, "segment_names")
  output_seg_prop = file.path(output_meshes, "segment_properties")
  # save meshes
  if (!dir.exists(output_path))
    dir.create(output_path)
  if (!dir.exists(output_meshes))
    dir.create(output_meshes)
  idx <- start_index
  for (msh in mesh_list){
    save_single_res_mesh(msh, output_path = output_meshes, name = as.character(idx))
    save_single_res_manifest(output_path = output_meshes, name = as.character(idx))
    idx <- idx + 1
  }

  # save annotations
  if (!dir.exists(output_seg_names))
    dir.create(output_seg_names)
  if (!dir.exists(output_seg_prop))
    dir.create(output_seg_prop)
  save_single_res_info(output_info)
  write_meshes_info(mesh_list,
                    start_idx = start_index,
                    output_seg_names = output_seg_names,
                    output_seg_prop = output_seg_prop)
}

# ---- Skeletons

save_skel_info <- function(output_info, radius = TRUE) {
  info_dict <- list(
    "@type" = "neuroglancer_skeletons",
    "transform": c(1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0),
    "segment_properties"= "seg_props"
  )
  if (isTRUE(radius))
    info_dict["vertex_attributes"] <- c(
      list("id"= "radius", "data_type"="float32", "num_components"=1)
    )

  jsonlite::write_json(
    info_dict,
    output_info,
    auto_unbox = TRUE,
  )
}

save_skel_props <- function(skels, output_skels, start_idx) {
  idcs <- start_idx:(start_idx+length(skels))
  properties <- list(
    "@type" = "neuroglancer_segment_properties",
    "inline" = list(
      "ids" = idcs,
      "properties" = c(
        lapply(idcs, function(x) list(id = idcs[[x]], type='label',
                                      values=names(skels)[[x]]))
      )
    )
  )
  jsonlite::write_json(
    properties,
    file.path(output_skels, "seg_props", "info"),
    auto_unbox = TRUE,
  )
}


save_single_res_skel <- function(skel, output_path, name) {
  vertex_positions <- xyzmatrix(skel)
  edges <- cbind(skel$d$PointNo[skel$d$Parent!=-1],
                 skel$d$Parent[skel$d$Parent!=-1])

  fc = file(file.path(output_path, name), "wb")
  writeBin(as.integer(dim(vertex_positions)[[1]]), fc)
  writeBin(as.integer(floor(dim(edges)[[1]]/2)), fc)
  writeBin(as.vector(t(vertex_positions)), fc, size = 4L)
  writeBin(as.integer(t(edges)), fc, size = 4L)
  close(fc)
}

read_single_res_mesh <- function(filename) {
  # read from binary
  fc <- file(filename, "rb")
  num_nodes <- readBin(fc, integer(), n = 1, size = 4L)
  nodes_pos = readBin(fc, double(), n = 3 * num_nodes, size = 4L)
  nodes_pos = matrix(nodes_pos, ncol = 3, byrow=T)
  ee <- readBin(fc, integer(), n = 1024, size = 4L)
  edges <- c(ee)
  while (length(ee) > 0) { # TODO optimize
    ee <- readBin(fc, integer(), n = 1024, size = 4L)
    edges <- c(edges, ee)
  }
  edges = matrix(edges, ncol = 3, byrow=T)
  edges = edges + 1
  close(fc)
  # create mesh
  mesh <- list(vb = vertices <- t(cbind(nodes_pos, 1)),
               it = t(edges),
               primitivetype = "triangle")
  class(mesh) <- c("mesh3d", "shape3d")
  mesh
}

read_skeleton <- function(filename) {
  # read from binary
  fc <- file(filename, "rb")
  num_nodes <- readBin(fc, integer(), n = 1, size = 4L)
  num_edges <- readBin(fc, integer(), n = 1, size = 4L)
  nodes_pos = readBin(fc, double(), n = 3 * num_nodes, size = 4L)
  nodes_pos = matrix(nodes_pos, ncol = 3, byrow=T)
  edges = readBin(fc, integer(), n = 2 * num_edges, size = 4L)
  edges = matrix(edges, ncol = 2, byrow=T)
  edges = edges + 1
  close(fc)
  # create a neuron
  d <- data.frame(
    PointNo = 1:num_nodes,
    X = nodes_pos[,1],
    Y = nodes_pos[,2],
    Z = nodes_pos[,3]
  )
  parent_mapping <- as.list(edges[,2])
  names(parent_mapping) <- as.character(edges[,1])
  d$Parent <- sapply(1:num_nodes, function(x) {
    m <- parent_mapping[[as.character(x)]]
    if (is.null(m)) -1 else m
  })
  as.neuron(d)
}

read_skeletons <- function(in_path) {
  if (!dir.exists(in_path))
    stop("This folder doesn't exist!")
  sk_fns <- list.files(in_path, pattern = "[0-9]+")
  sk_paths <- sapply(sk_fns, function(x) file.path(in_path, x))
  nl <- as.neuronlist(lapply(pps, read_skeleton))
  #if (file.exists()) #TODO check neuron names from seg_prop
  nl
}