Compact sparse matrix

compact.py

#! /usr/bin/python

from __future__ import division
from __future__ import print_function

import csv
import argparse
import numpy
import numpy.linalg
import numpy.random
import scipy.spatial

# Parse command line arguments
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Nuclear Segmentation with Fiji")
    parser.add_argument('--input-csv')
    
    args = parser.parse_args()
    objectListFile = args.input_csv

# Read objects
reader = csv.reader(open(objectListFile))
voxels = []
for index, line in enumerate(reader):
    if index > 0:
        voxel = {}
        voxel['coords'] = numpy.array([[line[1], line[2], line[3]]])
        voxel['volume'] = line[4]
        voxel['values'] = numpy.array([[line[5], line[7]]])
        voxel['neighs'] = numpy.full((3, 3), -1, dtype=int)
        voxel['neighs'][1, 1] = index - 1
        voxels.append(voxel)

# Create coordinate matrix and object list
matrix = numpy.empty([len(voxels), 2])
list = numpy.arange(len(voxels))

for index in range(0, len(voxels)):
    matrix[index] = voxels[index]['coords'][0, 0:2:1]

# Compute distances
tree = scipy.spatial.cKDTree(matrix)
distances, neighbors = tree.query(matrix, k=25)

worklist = list.copy()
numpy.random.shuffle(worklist)

### Check if neighbor's spot is free
def can_reference(x, y, item, neigh, voxels):
    
    x1 = x + 1
    y1 = y + 1
    x2 = -x + 1
    y2 = -y + 1
    
    if voxels[item]['neighs'][x1, y1] != -1 \
        and voxels[item]['neighs'][x1, y1] != neigh:
        return False
    
    if voxels[neigh]['neighs'][x2, y2] != -1 \
        and voxels[neigh]['neighs'][x2, y2] != item:
        return False
    
    return True

### Check if neighbor's neighbors conflict with item's neighbors
def can_merge(Cx, Cy, item, neigh, voxels):
    
    updates = []
    
    for Dx in range(Cx - 1, Cx + 2):
        for Dy in range(Cy - 1, Cy + 2):
            if Dx in range(-1,2) and Dy in range(-1,2) \
                and not (Dx == 0 and Dy == 0) \
                and not (Dx == Cx and Dy == Cy):
                
                current = voxels[item]['neighs'][Dx + 1, Dy + 1]
                if current != -1:
                    Zx = -Dx + Cx
                    Zy = -Dy + Cy

                    if can_reference(Zx, Zy, current, neigh, voxels):
                        it = {}
                        it['x'] = Zx
                        it['y'] = Zy
                        it['item'] = current
                        it['neigh'] = neigh
                        updates.append(it)
                    else:
                        return False
    return updates

### Reference this neighbor
def reference_neighbor(x, y, item, neigh, voxels):
    
    x1 = x + 1
    y1 = y + 1
    x2 = -x + 1
    y2 = -y + 1
    
    voxels[item]['neighs'][x1, y1] = neigh
    voxels[neigh]['neighs'][x2, y2] = item

### Adopt neighbor's neighbors as our own
def update_references(items, voxels):
    
    for item in items:
        reference_neighbor(item['x'], item['y'], \
            item['item'], item['neigh'], voxels)


for item in worklist:
    neigh_count = 0
    for nord in range(1, 25):
        if (neigh_count == 8):
            break
        neigh = neighbors[item][nord]
        A = matrix[item]
        B = matrix[neigh]
        BA = B - A
        C = (numpy.round(BA / numpy.linalg.norm(BA))).astype(int)
        Cx = C[0]
        Cy = C[1]

        if can_reference(Cx, Cy, item, neigh, voxels):
            refs_item = can_merge(Cx, Cy, item, neigh, voxels)
            refs_neigh = can_merge(-Cx, -Cy, neigh, item, voxels)
            if refs_item != False and refs_neigh != False:
                reference_neighbor(Cx, Cy, item, neigh, voxels)
                update_references(refs_item, voxels)
                update_references(refs_neigh, voxels)
                neigh_count = neigh_count + 1