Scripts: capmask.py

# -----------------------------------------------------------------------------
# Try to automatically mask cylindrical capillary from x-ray tomogram.
# Operates on highlighted volume in volume dialog when this script is opened.
# Creates a new volume data set.
#
# Tested with Chimera 1.5, July 19, 2010 daily build.
#
# Method: Look at two xy planes at 10% and 90% z values.  Sum along y axis
#   to get a 1d funtion.  Take first maxima on each end that are above
#   (mean + 2 standard deviations) as ellipse cross-section boundaries.
#   Likewise along x axis. Connect two circles to find cylinder axis and radius.
#
def fit_and_mask_cylinder(zpos = (0.10, 0.90), num_sd = 2.0,
                          scale_radius = 0.95, scale_height = 1.05,
                          color = (0,0,.7,1), subdivisions = 100,
                          mask = True):

    from VolumeViewer import active_volume
    v = active_volume()                                 # Highlighted in dialog.
    center, axis, radius, height = fit_cylinder(v.data, zpos, num_sd)
    radius *= scale_radius                              # Scale cylinder radius.
    height *= scale_height                              # Scale cylinder height.
    s = show_cylinder(center, axis, radius, height, subdivisions, color, v)
    if mask:
        from chimera import runCommand
        runCommand('mask #%d #%d' % (v.id, s.id))       # Mask volume.

# -----------------------------------------------------------------------------
# Returns center, axis, radius, height in volume data index units.
#
def fit_cylinder(data, zpos, num_sd):

    # Get data planes.
    xsize, ysize, zsize = data.size
    z1, z2 = [int(round(zp*zsize)) for zp in zpos]
    m1, m2 = data_z_plane(data,z1), data_z_plane(data,z2)

    # Set values below threshold to 0.
    m1.put(m1 < m1.mean() + num_sd * m1.std(), 0)
    m2.put(m2 < m2.mean() + num_sd * m2.std(), 0)

    # Find box bounding cylinder using 2 z-planes.
    from math import sqrt
    sdx = num_sd / sqrt(ysize)
    sdy = num_sd / sqrt(xsize)
    cx1, rx1 = local_maxima_bounds(m1.sum(axis=0), sdx)
    cy1, ry1 = local_maxima_bounds(m1.sum(axis=1), sdy)
    cx2, rx2 = local_maxima_bounds(m2.sum(axis=0), sdx)
    cy2, ry2 = local_maxima_bounds(m2.sum(axis=1), sdy)
    if cx1 is None or cy1 is None or cx2 is None or cy2 is None:
        print 'Projection with no local maxima above threshold'
        return

    # Correct radius for axis tilted away from z.
    dx,dy,dz = (cx2-cx1, cy2-cy1, z2-z1)
    from math import sqrt
    cosx = dz/sqrt(dx*dx+dz*dz)
    cosy = dz/sqrt(dy*dy+dz*dz)
    # Four radius values from x and y widths on z1 and z2 planes.
    radii = (rx1*cosx, rx2*cosx, ry1*cosy, ry2*cosy)
    radius = sum(radii)/len(radii)  # Use average radius.

    # Calculate cylinder center, axis and height.
    center = (0.5*(cx1+cx2), 0.5*(cy1+cy2), 0.5*(z1+z2))
    l = sqrt(dx*dx+dy*dy+dz*dz)
    axis = (dx/l,dy/l,dz/l)
    height = l * float(zsize)/(z2-z1)

    return center, axis, radius, height

# -----------------------------------------------------------------------------
#
def local_maxima_bounds(data, num_sd):

    cutoff = data.mean() + num_sd * data.std()
    local_max = ((data[1:-1] >= cutoff) &
                 (data[2:] <= data[1:-1]) &
                 (data[:-2] <= data[1:-1]))
    mi = local_max.nonzero()[0]           # indices of local maxima above mean
    if len(mi) == 0:
        return None, None

    i0 = mi[0]+1
    i1 = mi[-1]+1
    c = 0.5 * (i0 + i1)
    r = 0.5 * (i1 - i0)
    return c, r

# -----------------------------------------------------------------------------
#
def data_z_plane(data, z):

    xsize, ysize, zsize = data.size
    m3d = data.matrix((0,0,z), (xsize,ysize,1))
    m = m3d[0,:,:].copy()
    return m

# -----------------------------------------------------------------------------
#
def show_cylinder(center, axis, radius, height, divisions, color, volume):

    transform = volume.data.ijk_to_xyz_transform
    varray, tarray = cylinder_geometry(center, axis, radius, height,
                                       divisions, transform)
    name = volume.name + ' cyl'
    align = volume.openState.xform
    s = create_cylinder_model(varray, tarray, color, align, name)

    return s

# -----------------------------------------------------------------------------
# Mesh for cylinder.
#
def cylinder_geometry(center, axis, radius, height, divisions, transform):

    from math import pi
    circumf = 2*pi*radius
    if circumf > height:
        nc = divisions
        nz = max(1, int(divisions * height / circumf))
    else:
        nc = max(6, int(divisions * circumf / height))
        nz = divisions

    # Surface for cylinder with axis along z.
    from Shape.shapecmd import cylinder_geometry
    varray, tarray = cylinder_geometry(radius, height, nz, nc, caps = True)

    # Tilt cylinder axis
    import Matrix
    xa,ya,za = Matrix.orthonormal_frame(axis)
    cx,cy,cz = center
    tf = ((xa[0], ya[0], za[0], cx),
          (xa[1], ya[1], za[1], cy),
          (xa[2], ya[2], za[2], cz))
    from _contour import affine_transform_vertices
    affine_transform_vertices(varray, tf)
    affine_transform_vertices(varray, transform)

    return varray, tarray

# -----------------------------------------------------------------------------
#
def create_cylinder_model(vertices, triangles, color, align, name):

    # Create surface model
    import _surface
    s = _surface.SurfaceModel()
    s.name = name
    p = s.addPiece(vertices, triangles, color)
    p.displayStyle = p.Mesh

    # Show model
    from chimera import openModels
    openModels.add([s])
    if align:
        s.openState.xform = align
    return s

# -----------------------------------------------------------------------------
#
fit_and_mask_cylinder()