Generate HeightField

#!BPY
# coding: utf-8
"""
Name: '********** Generate Heightfield ********** '
Blender: 249
Group: 'Export'
Tooltip: 'Generate Heightfield'
"""

__author__ = ["Ricardo Marmolejo García"]
__url__ = ("http://es.linkedin.com/in/ricardomarmolejogarcia", "https://twitter.com/makiolo", "http://http://blogricardo.wordpress.com/")
__version__ = "1.00"
__bpydoc__ = """\
Ricardo Marmolejo García
"""

import Blender
from Blender import *
import bpy
import math
from BPyMesh import getMeshFromObject
from BPyObject import getDerivedObjects
import Image
import PIL

# documentacion
# www.blender.org/documentation/249PythonDoc/index.html

############################ VARIABLES CONFIGURABLES #############

# ruta de la imagen con la información de altura
heightfield = "/home/makiolo/Dropbox/heightfield/heightmap200.png"

# dimensiones del heightfield en unidades de juego (metro)
width = 1000
depth = 1000
height = 100

############################ AUXILIARES - NO TOCAR ############################

# hay que generar los vertex e indices del mesh
vertex = []
indexs = []

# datos de la imagen
data = {}

# numero de vertices en ancho y fondo
width_verts = 0
depth_verts = 0

# numero de tiles
width_tiles = 0
depth_tiles = 0

# ancho y fondo de cada tile
width_step = 0
depth_step = 0

#################################################################################

# establecimiendo de globales mediante el tamaño de la imagen
def setSizeInfo(_width, _height):

    global width_verts, depth_verts, width_tiles, depth_tiles, width_step, depth_step

    # numero de vertices en ancho y fondo
    width_verts = _width
    depth_verts = _height

    # numero de tiles
    width_tiles = width_verts - 1
    depth_tiles = depth_verts - 1

    # ancho y fondo de cada tile
    width_step = width / width_tiles
    depth_step = depth / depth_tiles

# funcion que determina la altura de cada vertice
def getFunction(x, z):

    # La altura viene dada por la imagen
    # aunque podría ser dada por una función matemática
    y = data["%d-%d" % (x, z)]

    #y = math.sin(0.3 * x) * math.cos(0.1 * z)
    return y

class VertexData:
    def __init__(self, pos, uv):
        self.pos = pos
        self.uv = uv

    def __repr__(self):
        return ("Posicion%s - UV%s" % (self.pos, self.uv))

class Vector2:
    def __init__(self, u, v):
        self.u = u
        self.v = v

    def __repr__(self):
        return ("(%.2f, %.2f)" % (self.u, self.v))

class Vector3:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        return ("(%.2f, %.2f, %.2f)" % (self.x, self.y, self.z))

class Triangle:
    def __init__(self, idx1, idx2, idx3):
        self.idx1 = idx1
        self.idx2 = idx2
        self.idx3 = idx3

    def __repr__(self):
        return ("Triangle (%d, %d, %d)" % (self.idx1, self.idx2, self.idx3))

def generate_heightfield():
    editmode = Window.EditMode()    # are we in edit mode?  If so ...
    if editmode: Window.EditMode(0) # leave edit mode before getting the mesh

    Blender.Window.DrawProgressBar(0.0, 'Generando heightfield.')

    # generar vertex
    for c in range(width_verts):
        for f in range(depth_verts):
            vertex.append( VertexData(Vector3(  f * width_step,     # x
                                getFunction(f, c),                  # y
                                c * depth_step),                    # z
                                Vector2(float(f) / width_tiles,             # coord u
                                float(c) / depth_tiles)                     # coord v
                        ))

    Blender.Window.DrawProgressBar(0.2, 'Generando heightfield.')

    # calcular maximo y minimo en altura
    # se calcula la maxima amplitud
    # normalizamos y escalamos a la altura solicitada
    max_height = -9999999
    min_height = +9999999

    for v in vertex:
        if v.pos.y < min_height:
            min_height = v.pos.y

        if v.pos.y > max_height:
            max_height = v.pos.y

    distancia = max_height - min_height

    for v in vertex:
        v.pos.y = (v.pos.y / distancia) * height

    Blender.Window.DrawProgressBar(0.4, 'Generando heightfield.')

    # calcular centro del mesh
    center = Vector3(0, 0, 0)
    count = 0

    for v in vertex:
        center.x += v.pos.x
        center.y += v.pos.y
        center.z += v.pos.z
        count = count + 1

    if count > 0:
        center.x /= count
        center.y /= count
        center.z /= count

    Blender.Window.DrawProgressBar(0.6, 'Generando heightfield.')

    # centrar pivote
    for v in vertex:
        v.pos.x -= center.x
        v.pos.y -= center.y
        v.pos.z -= center.z

    Blender.Window.DrawProgressBar(0.7, 'Generando heightfield.')

    # generar indices
    # A
    #     i + j             ___         i + j + width_verts
    #                       | /|
    #     i + j + 1         |/_|        i + j + width_verts + 1
    #
    #
    # B
    #     i + j             ___         i + j + width_verts
    #                       |\ |
    #     i + j + 1         |_\|        i + j + width_verts + 1
    #

    (CLOCKWISE_A, CLOCKWISE_B,
    ANTICLOCKWISE_A, ANTICLOCKWISE_B)=([ "%d" % i for i in range(4) ])

    mode = CLOCKWISE_A

    i = j = 0
    for c in range(width_tiles): # z
        for f in range(depth_tiles): # x

            if mode == CLOCKWISE_A:

                indexs.append(Triangle(i + j,                          i + j + width_verts,          i + j + 1))
                indexs.append(Triangle(i + j + width_verts,            i + j + width_verts + 1,      i + j + 1))

            elif mode == CLOCKWISE_B:

                indexs.append(Triangle(i + j,                          i + j + width_verts + 1,      i + j + 1))
                indexs.append(Triangle(i + j,                          i + j + width_verts,          i + j + width_verts + 1))

            if mode == ANTICLOCKWISE_A:

                indexs.append(Triangle(i + j + 1,                           i + j + width_verts,          i + j))
                indexs.append(Triangle(i + j + 1,                           i + j + width_verts + 1,      i + j + width_verts))

            elif mode == ANTICLOCKWISE_B:

                indexs.append(Triangle(i + j + 1,                           i + j + width_verts + 1,      i + j))
                indexs.append(Triangle(i + j + width_verts + 1,             i + j + width_verts,          i + j))

            i = i + 1

        j = j + 1

    Blender.Window.DrawProgressBar(0.9, 'Generando heightfield.')

    # mostrar resultados
    coords = []
    faces = []

    for v in vertex:
        coords.append( [v.pos.x, v.pos.y, v.pos.z] )

    for i in indexs:
        faces.append( [i.idx1, i.idx2, i.idx3] )

    me = bpy.data.meshes.new('myMesh')          # create a new mesh

    me.addUVLayer("coord_uv")

    me.verts.extend(coords)
    me.faces.extend(faces)

    '''
    TODO:
    Establecer colores a los vertices
    Tengo calculadas las coordenadas UV pero no son asginadas

    me.vertexColors = 1
    me.faces[1].col[0].r = 255
    me.faces[1].col[1].g = 255
    me.faces[1].col[2].b = 255
    '''

    scn = bpy.data.scenes.active
    ob = scn.objects.new(me, 'myHeightfield')

    me.update()
    if editmode: Window.EditMode(1)
    Blender.Redraw()

    Blender.Window.DrawProgressBar(1.0, 'Heightfield generado.')

    salida = "Heightfield generado\n"
    Blender.Draw.PupMenu(salida)

def leer_imagen():
    # lee la imagen con la informacion de altura
    im = Image.open(heightfield)
    #print im.format, im.size, im.mode
    # Abre la imagen con el visor del SO
    #im.show()

    _width = im.size[0]
    _height = im.size[1]

    setSizeInfo(_width, _height)

    # grises 8 bits
    if(im.mode == "L"):

        for y in range(_height):
            for x in range(_width):
                pixel = im.getpixel( (x,y) )
                data["%d-%d" % (x,y)] = float(pixel)

        return True

    # RGB 32 bits
    elif(im.mode == "RGB"):

        for y in range(_height):
            for x in range(_width):
                pixel = im.getpixel( (x,y) )
                data["%d-%d" % (x,y)] = float(pixel[0] + pixel[1] + pixel[2]) / 3.0

        return True

    else:
        print("Formato invalido")
        return False

if __name__=='__main__':
    if leer_imagen():
        generate_heightfield()
    else:
        print("No se ha generado el mesh")