render_tools.py

import pygame
from pygame.locals import *

from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GL import shaders
import numpy as np

sphereQuadric = None
   
def initLights():
    glLightfv(GL_LIGHT0, GL_AMBIENT,  [0.2, 0.2, 0.2, 1])
    glLightfv(GL_LIGHT0, GL_DIFFUSE,  [1,1,1,1])
    glLightfv(GL_LIGHT0, GL_SPECULAR, [1,1,1,1])
    glLightfv(GL_LIGHT0, GL_POSITION, [3,-10,-10,0])
    glEnable(GL_LIGHT0)

    glLightfv(GL_LIGHT1, GL_AMBIENT,  [0.2, 0.2, 0.2, 1])
    glLightfv(GL_LIGHT1, GL_DIFFUSE,  [1,1,1,1])
    glLightfv(GL_LIGHT1, GL_SPECULAR, [1,1,1,1])
    glLightfv(GL_LIGHT1, GL_POSITION, [3,10,-10,0])
    glEnable(GL_LIGHT1)

    glLightfv(GL_LIGHT2, GL_AMBIENT,  [0.2, 0.2, 0.2, 1])
    glLightfv(GL_LIGHT2, GL_DIFFUSE,  [1,1,1,1])
    glLightfv(GL_LIGHT2, GL_SPECULAR, [1,1,1,1])
    glLightfv(GL_LIGHT2, GL_POSITION, [3,-10,10,0])
    glEnable(GL_LIGHT2)
    
    glLightfv(GL_LIGHT3, GL_AMBIENT,  [0.2, 0.2, 0.2, 1])
    glLightfv(GL_LIGHT3, GL_DIFFUSE,  [1,1,1,1])
    glLightfv(GL_LIGHT3, GL_SPECULAR, [1,1,1,1])
    glLightfv(GL_LIGHT3, GL_POSITION, [3,10,10,0])
    glEnable(GL_LIGHT3)

    glEnable(GL_LIGHTING)
    glLightModelf(GL_LIGHT_MODEL_LOCAL_VIEWER, GL_TRUE)
    glLightModeli(GL_LIGHT_MODEL_TWO_SIDE, GL_FALSE)
   
def drawMesh(pd, mesh, offset, color, alpha=1.0):
    global shader
    global UNIFORM_LOCATIONS
    shaders.glUseProgram(shader)

    mv = glGetFloatv(GL_MODELVIEW_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS["modelview_matrix"], 1, GL_FALSE, mv)
    pm = glGetFloatv(GL_PROJECTION_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS["projection_matrix"], 1, GL_FALSE, pm)
    glUniform3fv(UNIFORM_LOCATIONS["surface_color"], 1, color)
    glUniform1f(UNIFORM_LOCATIONS["alpha"], alpha)
    glUniform1f(UNIFORM_LOCATIONS["shininess"], 5.0)
    glUniform1f(UNIFORM_LOCATIONS["specular_factor"], 0.2)
     
    faces = mesh.getFaces()

    #glMaterialfv(GL_FRONT_AND_BACK, GL_AMBIENT, color)
    #glMaterialfv(GL_FRONT_AND_BACK, GL_DIFFUSE, color)
    #glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR, [1,1,1,1])
    #glMaterialf(GL_FRONT_AND_BACK, GL_SHININESS, 15.0)
    #glColor3fv(color)

    v = np.array(pd.getVertices())
    vertices = np.reshape(v[offset:], -1)
    normals = np.reshape(np.negative(np.array(mesh.getVertexNormals())), -1)

    #glEnableClientState(GL_VERTEX_ARRAY)
    #glEnableClientState(GL_NORMAL_ARRAY)
    #glVertexPointer(3, GL_DOUBLE, 0, vertices)
    #glNormalPointer(GL_DOUBLE, 0, normals)
    
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, vertices);
    glEnableVertexAttribArray(1);
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, normals);

    glDrawElements(GL_TRIANGLES, 3 * mesh.numFaces(), GL_UNSIGNED_INT, mesh.getFaces())

    #glDisableClientState(GL_VERTEX_ARRAY)
    #glDisableClientState(GL_NORMAL_ARRAY)
    glDisableVertexAttribArray(0)
    glDisableVertexAttribArray(2)
    shaders.glUseProgram(0)

def drawMeshWithTexture(pd, mesh, offset, texture_id, tex_coords, alpha=1.0, use_texture=1.0):
    """
    使用纹理渲染mesh。
    
    Args:
        pd: ParticleData对象
        mesh: Mesh对象
        offset: 顶点偏移量
        texture_id: OpenGL纹理ID
        tex_coords: np.array(N, 2) 纹理坐标数组
        alpha: 透明度
        use_texture: 是否使用纹理（1.0=使用，0.0=不使用）
    """
    global shader_textured
    global UNIFORM_LOCATIONS_TEXTURED
    
    if shader_textured is None:
        print("Warning: Textured shader not initialized, falling back to regular drawMesh")
        drawMesh(pd, mesh, offset, [1.0, 0.0, 0.0], alpha)
        return
    
    shaders.glUseProgram(shader_textured)

    mv = glGetFloatv(GL_MODELVIEW_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS_TEXTURED["modelview_matrix"], 1, GL_FALSE, mv)
    pm = glGetFloatv(GL_PROJECTION_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS_TEXTURED["projection_matrix"], 1, GL_FALSE, pm)
    glUniform1f(UNIFORM_LOCATIONS_TEXTURED["alpha"], alpha)
    glUniform1f(UNIFORM_LOCATIONS_TEXTURED["shininess"], 5.0)
    glUniform1f(UNIFORM_LOCATIONS_TEXTURED["specular_factor"], 0.2)
    glUniform1f(UNIFORM_LOCATIONS_TEXTURED["use_texture"], use_texture)
    
    # 绑定纹理
    if texture_id is not None and use_texture > 0.5:
        glActiveTexture(GL_TEXTURE0)
        glBindTexture(GL_TEXTURE_2D, texture_id)
        glUniform1i(UNIFORM_LOCATIONS_TEXTURED["texture_sampler"], 0)
    
    faces = mesh.getFaces()

    v = np.array(pd.getVertices())
    vertices = np.reshape(v[offset:], -1)
    normals = np.reshape(np.negative(np.array(mesh.getVertexNormals())), -1)
    
    # 确保tex_coords是正确的格式
    if tex_coords is not None:
        tex_coords_flat = np.reshape(tex_coords, -1).astype(np.float32)
    else:
        # 如果没有提供纹理坐标，使用默认值
        num_vertices = len(v) - offset
        tex_coords_flat = np.zeros(num_vertices * 2, dtype=np.float32)
    
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, vertices)
    glEnableVertexAttribArray(1)
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, normals)
    glEnableVertexAttribArray(2)
    glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, 0, tex_coords_flat)

    glDrawElements(GL_TRIANGLES, 3 * mesh.numFaces(), GL_UNSIGNED_INT, mesh.getFaces())

    glDisableVertexAttribArray(0)
    glDisableVertexAttribArray(1)
    glDisableVertexAttribArray(2)
    
    # 解绑纹理
    if texture_id is not None:
        glBindTexture(GL_TEXTURE_2D, 0)
    
    shaders.glUseProgram(0)

def drawCustomMesh(vertices, faces, normals, color, alpha=1.0):
    """
    使用shader渲染自定义顶点和面的网格（优化版本）。
    
    Args:
        vertices: list or numpy array of shape (N, 3) - 顶点位置
        faces: numpy array of shape (M*3,) - 预计算的扁平面索引数组，或 list of 3-tuples（向后兼容）
        normals: numpy array of shape (N, 3) - 预计算的顶点法线（如果为None则每帧计算，性能低）
        color: [r, g, b] - 颜色
        alpha: float - 透明度
    """
    global shader
    global UNIFORM_LOCATIONS
    
    if len(vertices) == 0:
        return
    
    shaders.glUseProgram(shader)
    
    mv = glGetFloatv(GL_MODELVIEW_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS["modelview_matrix"], 1, GL_FALSE, mv)
    pm = glGetFloatv(GL_PROJECTION_MATRIX)
    glUniformMatrix4fv(UNIFORM_LOCATIONS["projection_matrix"], 1, GL_FALSE, pm)
    glUniform3fv(UNIFORM_LOCATIONS["surface_color"], 1, color)
    glUniform1f(UNIFORM_LOCATIONS["alpha"], alpha)
    glUniform1f(UNIFORM_LOCATIONS["shininess"], 5.0)
    glUniform1f(UNIFORM_LOCATIONS["specular_factor"], 0.2)
    
    # Convert vertices to numpy array (如果已经是numpy array则直接使用)
    if isinstance(vertices, list):
        vertices_np = np.array(vertices, dtype=np.float32)
    else:
        vertices_np = np.asarray(vertices, dtype=np.float32)
    
    # Flatten vertices array
    vertices_flat = vertices_np.flatten()
    
    # Handle normals (优先使用预计算的，否则每帧计算)
    if normals is not None and len(normals) > 0:
        # Use precomputed normals (fast path)
        normals_np = np.asarray(normals, dtype=np.float32)
        normals_flat = normals_np.flatten()
    else:
        # Fallback: compute normals on the fly (slow, for backward compatibility)
        if isinstance(faces, list):
            # Convert list to numpy array
            faces_np = np.array(faces)
            normals_np = np.zeros_like(vertices_np)
            # Compute face normals and accumulate to vertices
            for face in faces:
                if len(face) >= 3:
                    v0_idx, v1_idx, v2_idx = face[0], face[1], face[2]
                    if v0_idx < len(vertices_np) and v1_idx < len(vertices_np) and v2_idx < len(vertices_np):
                        v0 = vertices_np[v0_idx]
                        v1 = vertices_np[v1_idx]
                        v2 = vertices_np[v2_idx]
                        edge1 = v1 - v0
                        edge2 = v2 - v0
                        face_normal = np.cross(edge1, edge2)
                        face_normal_len = np.linalg.norm(face_normal)
                        if face_normal_len > 1e-6:
                            face_normal = face_normal / face_normal_len
                            normals_np[v0_idx] += face_normal
                            normals_np[v1_idx] += face_normal
                            normals_np[v2_idx] += face_normal
            norms = np.linalg.norm(normals_np, axis=1, keepdims=True)
            norms[norms < 1e-6] = 1.0
            normals_np = normals_np / norms
            normals_flat = normals_np.flatten()
        else:
            # If faces is already a flat array, we can't compute normals easily
            normals_np = np.zeros_like(vertices_np)
            normals_flat = normals_np.flatten()
    
    # Handle faces array (优先使用预计算的扁平数组)
    if isinstance(faces, np.ndarray) and faces.ndim == 1:
        # Precomputed flat array (fast path)
        faces_np = np.asarray(faces, dtype=np.uint32)
    elif isinstance(faces, list):
        # Convert list to flat array (slow, for backward compatibility)
        faces_array = []
        for face in faces:
            if len(face) >= 3:
                faces_array.extend([face[0], face[1], face[2]])
        faces_np = np.array(faces_array, dtype=np.uint32)
    else:
        # Already numpy array but might be 2D, flatten it
        faces_np = np.asarray(faces, dtype=np.uint32).flatten()
    
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, vertices_flat)
    glEnableVertexAttribArray(1)
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 0, normals_flat)
    
    glDrawElements(GL_TRIANGLES, len(faces_np), GL_UNSIGNED_INT, faces_np)
    
    glDisableVertexAttribArray(0)
    glDisableVertexAttribArray(1)
    shaders.glUseProgram(0)
    
def drawSphere(translation, radius, color, subDivision=16):
    global sphereQuadric
    speccolor = [1.0, 1.0, 1.0, 1.0]
    glMaterialfv (GL_FRONT_AND_BACK, GL_AMBIENT, color)
    glMaterialfv (GL_FRONT_AND_BACK, GL_DIFFUSE, color)
    glMaterialfv (GL_FRONT_AND_BACK, GL_SPECULAR, speccolor)
    glMaterialf (GL_FRONT_AND_BACK, GL_SHININESS, 5.0);
    glColor3fv(color)

    if (sphereQuadric == None):
        sphereQuadric = gluNewQuadric();
        gluQuadricNormals(sphereQuadric, GLU_SMOOTH);

    glPushMatrix ()
    glTranslated ((translation)[0], (translation)[1], (translation)[2])
    gluSphere(sphereQuadric, radius, subDivision, subDivision)
    glPopMatrix()

 
 
def drawText(position, textString, size = 36):     
    font = pygame.font.Font (None, size)
    textSurface = font.render(textString, True, (255,255,255,255), (0,0,0,255))     
    textData = pygame.image.tostring(textSurface, "RGBA", True)     
    glRasterPos2d(*position)     
    glDrawPixels(textSurface.get_width(), textSurface.get_height(), GL_RGBA, GL_UNSIGNED_BYTE, textData)    
  
def initShader():
    global shader
    global UNIFORM_LOCATIONS

    VERTEX_SHADER = shaders.compileShader("""#version 330
        uniform mat4 modelview_matrix;
        uniform mat4 projection_matrix;

        layout(location = 0) in vec3 position;
        layout(location = 1) in vec3 normal;

        out VertexData
        {
           vec3 normal;
           vec3 view;
        }
        outData;

        void main()
        {
            vec4 v_position = projection_matrix * modelview_matrix * vec4(position, 1.0);
            outData.normal = mat3(modelview_matrix) * normal;
            outData.view = -v_position.xyz;
            gl_Position = v_position;
        }
        """, GL_VERTEX_SHADER)
        
    FRAGMENT_SHADER = shaders.compileShader("""#version 330
        uniform vec3 surface_color;
        uniform float shininess;
        uniform float specular_factor;
        uniform float alpha;

        layout(location = 0, index = 0) out vec4 frag_color;

        in VertexData
        {
            vec3 normal;
            vec3 view;
        }
        inData;

        void main()
        {
            vec3 normal = inData.normal;

            if (!gl_FrontFacing)
            {
                normal = -inData.normal;
            }
            
            const vec3 ambient = vec3(0.6, 0.6, 0.6);
            const vec3 diffuse = vec3(1.0, 1.0, 1.0);
            const vec3 specular = vec3(1.0, 1.0, 1.0);	
            
            vec3 eye_n = normalize(normal);
            vec3 eye_v = normalize(inData.view);
            const vec3 eye_l = vec3(0.0, 0.0, 1.0);	
            float cos_ln = max(dot(eye_l, eye_n), 0.0);    
            vec3 h = normalize(eye_l + eye_v);
            float cos_nh = max(dot(eye_n, h), 0.0);   
            
            vec3 color = surface_color * (ambient + diffuse * cos_ln + specular_factor * specular * pow(cos_nh, shininess));

            frag_color = vec4(color, alpha);
        }
        """, GL_FRAGMENT_SHADER)
    shader = shaders.compileProgram(VERTEX_SHADER,FRAGMENT_SHADER)
    
    UNIFORM_LOCATIONS = {
            'modelview_matrix': glGetUniformLocation(shader, 'modelview_matrix'),
            'projection_matrix': glGetUniformLocation(shader, 'projection_matrix'),
            'surface_color': glGetUniformLocation(shader, 'surface_color'),
            'shininess': glGetUniformLocation(shader, 'shininess'),
            'specular_factor': glGetUniformLocation(shader, 'specular_factor'),
            'alpha': glGetUniformLocation(shader, 'alpha'),
        }

# 支持纹理的shader
shader_textured = None
UNIFORM_LOCATIONS_TEXTURED = None

def initTexturedShader():
    """
    初始化支持纹理的shader，用于将背景图片作为mesh表面纹理。
    """
    global shader_textured
    global UNIFORM_LOCATIONS_TEXTURED

    VERTEX_SHADER_TEXTURED = shaders.compileShader("""#version 330
        uniform mat4 modelview_matrix;
        uniform mat4 projection_matrix;

        layout(location = 0) in vec3 position;
        layout(location = 1) in vec3 normal;
        layout(location = 2) in vec2 texCoord;

        out VertexData
        {
           vec3 normal;
           vec3 view;
           vec2 texCoord;
        }
        outData;

        void main()
        {
            vec4 v_position = projection_matrix * modelview_matrix * vec4(position, 1.0);
            outData.normal = mat3(modelview_matrix) * normal;
            outData.view = -v_position.xyz;
            outData.texCoord = texCoord;
            gl_Position = v_position;
        }
        """, GL_VERTEX_SHADER)
        
    FRAGMENT_SHADER_TEXTURED = shaders.compileShader("""#version 330
        uniform sampler2D texture_sampler;
        uniform float shininess;
        uniform float specular_factor;
        uniform float alpha;
        uniform float use_texture;  // 0.0 = 不使用纹理, 1.0 = 使用纹理

        layout(location = 0, index = 0) out vec4 frag_color;

        in VertexData
        {
            vec3 normal;
            vec3 view;
            vec2 texCoord;
        }
        inData;

        void main()
        {
            vec3 normal = inData.normal;

            if (!gl_FrontFacing)
            {
                normal = -inData.normal;
            }
            
            const vec3 ambient = vec3(0.6, 0.6, 0.6);
            const vec3 diffuse = vec3(1.0, 1.0, 1.0);
            const vec3 specular = vec3(1.0, 1.0, 1.0);	
            
            vec3 eye_n = normalize(normal);
            vec3 eye_v = normalize(inData.view);
            const vec3 eye_l = vec3(0.0, 0.0, 1.0);	
            float cos_ln = max(dot(eye_l, eye_n), 0.0);    
            vec3 h = normalize(eye_l + eye_v);
            float cos_nh = max(dot(eye_n, h), 0.0);   
            
            // 计算光照
            float lighting = ambient.r + diffuse.r * cos_ln + specular_factor * specular.r * pow(cos_nh, shininess);
            
            // 根据use_texture决定使用纹理还是纯色
            vec3 color;
            if (use_texture > 0.5) {
                // 使用纹理
                vec3 texColor = texture(texture_sampler, inData.texCoord).rgb;
                color = texColor * lighting;
            } else {
                // 使用默认颜色（红色，用于肝脏）
                color = vec3(1.0, 0.0, 0.0) * lighting;
            }

            frag_color = vec4(color, alpha);
        }
        """, GL_FRAGMENT_SHADER)
    shader_textured = shaders.compileProgram(VERTEX_SHADER_TEXTURED, FRAGMENT_SHADER_TEXTURED)
    
    UNIFORM_LOCATIONS_TEXTURED = {
            'modelview_matrix': glGetUniformLocation(shader_textured, 'modelview_matrix'),
            'projection_matrix': glGetUniformLocation(shader_textured, 'projection_matrix'),
            'texture_sampler': glGetUniformLocation(shader_textured, 'texture_sampler'),
            'shininess': glGetUniformLocation(shader_textured, 'shininess'),
            'specular_factor': glGetUniformLocation(shader_textured, 'specular_factor'),
            'alpha': glGetUniformLocation(shader_textured, 'alpha'),
            'use_texture': glGetUniformLocation(shader_textured, 'use_texture'),
        }
        
def initGL(resX, resY):   
    pygame.init()
    display = (resX, resY)
    pygame.display.set_mode(display, DOUBLEBUF|OPENGL)

    glEnable (GL_DEPTH_TEST)
    glEnable (GL_NORMALIZE)
    glShadeModel (GL_SMOOTH)
    glPolygonMode (GL_FRONT_AND_BACK, GL_FILL)

    # Enable alpha blending for semi-transparent meshes
    glEnable(GL_BLEND)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
    
    gluPerspective(45, (display[0]/display[1]), 0.5, 50.0)
    
    initShader()
    initTexturedShader()  # 初始化纹理shader