crappy_swc_importer.py

crappy_swc_importer.py

import itertools
from collections import defaultdict
import bpy

def find_number_of_prelim_comments_lines(p):
    num = 0
    with open(p) as file:
        while next(file).startswith('#'):
            num += 1
            continue
    return num


def make_visual_dict(p):

    num_comment_lines = find_number_of_prelim_comments_lines(path)

    spamz = {}
    with open(p) as swc_file:
        for _ in range(num_comment_lines):
            next(swc_file)

        for line in swc_file:
            n, t, x, y, z, r, p = line.rstrip().split()
            spamz[int(n)] = int(t), float(x), float(y), float(z), float(r), int(p)
    return spamz


path = r'C:\Users\zeffi\Downloads\rubber\neuron_nmo\bergstrom\CNG version\291-1.CNG.swc'
points_dict = make_visual_dict(path)


locs = []
radii = []
edges = []
add_loc = locs.append
add_radii = radii.append
add_edge = edges.append
knot_dict = defaultdict(int)
for k, v in sorted(points_dict.items()):
    add_loc((v[1], v[2], v[3]))
    add_radii((v[4]))
    other_index = v[5]
    knot_dict[k-1] += 1
    if other_index > 0:
        add_edge((k-1, other_index-1))
        knot_dict[other_index-1] += 1

f_knots = [(v > 2) for k, v in knot_dict.items()]

mesh = bpy.data.meshes.new("Base_Profile_Data")
mesh.from_pydata(locs, edges, [])
mesh.update()

obj = bpy.data.objects.new("Base_Profile", mesh)

scene = bpy.context.scene
scene.objects.link(obj)
obj.select = True

a = obj.modifiers.new(type='SKIN', name='sv_skin')
b = obj.modifiers.new(type='SUBSURF', name='sv_subsurf')
a.branch_smoothing = 1.0

f_r = list(itertools.chain(*zip(radii, radii)))
mesh.skin_vertices[0].data.foreach_set('radius', f_r)
mesh.skin_vertices[0].data.foreach_set('use_loose', f_knots)