PDS IMG to 3D Mesh

main.py

from PIL import Image
from pyhull.delaunay import DelaunayTri
import numpy as np
import gdal
import operator
import struct
import sys

def load(path):
    print 'load'
    image = gdal.Open(path)
    data = image.ReadAsArray()
    data[data < -1e9] = np.nan
    return data

def normalize(data):
    lo = np.nanmin(data)
    hi = np.nanmax(data)
    pct = (data - lo) / (hi - lo)
    return pct

def array_to_png(data, path):
    print 'array_to_png'
    pct = normalize(data)
    pct = pct ** 0.5
    scaled = (255.0 * pct).astype(np.uint8)
    im = Image.fromarray(scaled)
    im.save(path)

def derivative(data):
    print 'derivative'
    diffs = []
    for dy in xrange(-1, 2):
        for dx in xrange(-1, 2):
            if dx and dy:
                continue
            if not dx and not dy:
                continue
            shifted = data
            shifted = np.roll(shifted, dy, axis=0)
            shifted = np.roll(shifted, dx, axis=1)
            diff = abs(data - shifted)
            diffs.append(diff)
    return reduce(np.maximum, diffs)

def select_points(data, threshold, step):
    print 'select_points'
    x, z = (data > threshold).nonzero()
    x = x[x % step == 0]
    z = z[z % step == 0]
    return zip(x, z)

def triangulate(original, data):
    print 'triangulate'
    data[np.isnan(data)] = 0
    print np.histogram(data)
    sets = [
        set(select_points(data, 0.9, 1)),
        set(select_points(data, 0.8, 1)),
        set(select_points(data, 0.7, 1)),
        set(select_points(data, 0.6, 1)),
        set(select_points(data, 0.5, 1)),
        set(select_points(data, 0.4, 2)),
        set(select_points(data, 0.3, 4)),
        set(select_points(data, 0.2, 8)),
        set(select_points(data, 0.1, 16)),
        set(select_points(data, 0.0, 32)),
    ]
    print [len(x) for x in sets]
    points = list(reduce(operator.or_, sets))
    print len(points)
    points = [(x, z) for x, z in points if not np.isnan(original[x][z])]
    print len(points)
    tri = DelaunayTri(points)
    points = [(x, original[x][z] * 2, z) for x, z in tri.points]
    positions = []
    for a, b, c in tri.vertices:
        positions.append(points[c])
        positions.append(points[b])
        positions.append(points[a])
    return positions

def save_binary_stl(positions, path):
    print 'save_binary_stl'
    p = positions
    data = []
    data.append('\x00' * 80)
    data.append(struct.pack('<I', len(p) / 3))
    for vertices in zip(p[::3], p[1::3], p[2::3]):
        data.append(struct.pack('<fff', 0.0, 0.0, 0.0))
        for vertex in vertices:
            data.append(struct.pack('<fff', *vertex))
        data.append(struct.pack('<H', 0))
    data = ''.join(data)
    with open(path, 'wb') as fp:
        fp.write(data)

def main():
    data = load(sys.argv[1])
    d0 = data
    d1 = derivative(d0)
    d2 = derivative(d1)
    pct = normalize(d2)
    pct = pct ** 0.5
    positions = triangulate(d0, pct)
    save_binary_stl(positions, 'output.stl')
    # array_to_png(d0, 'd0.png')
    # array_to_png(d1, 'd1.png')
    # array_to_png(d2, 'd2.png')

if __name__ == '__main__':
    main()