forked from hongbo-zhu-cn/Pymol-script-repo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
wfmesh.py
175 lines (141 loc) · 6.38 KB
/
wfmesh.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
'''
See more here: http://www.pymolwiki.org/index.php/wfmesh
###############################################
# File: WFObj.py
# Author: Dan Kulp
# Creation Date: 5/13/05
#
# Notes:
# Create openGL objects from a wavefront (obj) file
###############################################
'''
import os
import re
import math
from pymol.opengl.gl import *
from pymol.callback import Callback
from pymol import cmd
# Wrapper Function, to create a given WFObj with a specific name (flip = 1 if OpenFX + Crossroads used)
def createWFObj(file, name,translate=[0,0,0],flip=0):
obj = WFMesh(file,translate,flip)
cmd.load_callback(obj,name)
# Class for Wavefront Mesh
class WFMesh(Callback):
verts = [] # list of vertices
polys = [] # list of poylgons
pnorms = [] # list of polynomal normals
vnorms = {} # dict. of vertex normals
vavenorms = [] # list of vertex normals, redundant -- I'm far froma python pro.
sections = {} # list of sections of mesh
# Read mesh into memory
def readOBJ(self,file):
if os.path.exists(file):
input = open(file,'r')
for line in input:
dat = re.split("\s+", line)
# Find vertex line
if line[0] == 'v' and line[1] != 't' and line[1] != 'n': self.verts.append([dat[1],dat[2],dat[3]])
# Find polygon line
if line[0] == 'f': self.polys.append([dat[1],dat[2],dat[3]])
# Find section line
if line[0] == 'g': self.sections[len(self.polys)] = dat[1]
# Compute the normals for each polygon and each vertex
def computeNorms(self):
# Compute norms for each polygon
for p in self.polys:
v12 = [float(self.verts[int(p[1])-1][0]) - float(self.verts[int(p[0])-1][0]),\
float(self.verts[int(p[1])-1][1]) - float(self.verts[int(p[0])-1][1]),\
float(self.verts[int(p[1])-1][2]) - float(self.verts[int(p[0])-1][2]) \
]
v13 = [float(self.verts[int(p[2])-1][0]) - float(self.verts[int(p[0])-1][0]),\
float(self.verts[int(p[2])-1][1]) - float(self.verts[int(p[0])-1][1]),\
float(self.verts[int(p[2])-1][2]) - float(self.verts[int(p[0])-1][2]) \
]
# Compute poly normal
polynorm = self.cross(v12,v13)
norm = self.normalize(polynorm)
# Files created by OpenFX, Crossroads combination need have their normals flipped
if self.flip:
norm[0] = -norm[0]
norm[1] = -norm[1]
norm[2] = -norm[2]
# Append poly norm to polygonal norm array
self.pnorms.append(norm)
# Add norm to each vertexes norm..
try:
self.vnorms[int(p[0])-1] = [float(self.vnorms[int(p[0])-1][0]) + norm[0],
float(self.vnorms[int(p[0])-1][1]) + norm[1],
float(self.vnorms[int(p[0])-1][2]) + norm[2]
]
except:
self.vnorms[int(p[0])-1] = [norm[0],norm[1],norm[2]]
try:
self.vnorms[int(p[1])-1] = [float(self.vnorms[int(p[1])-1][0]) + norm[0],
float(self.vnorms[int(p[1])-1][1]) + norm[1],
float(self.vnorms[int(p[1])-1][2]) + norm[2]
]
except:
self.vnorms[int(p[1])-1] = [norm[0],norm[1],norm[2]]
try:
self.vnorms[int(p[2])-1] = [float(self.vnorms[int(p[1])-1][0]) + norm[0],
float(self.vnorms[int(p[1])-1][1]) + norm[1],
float(self.vnorms[int(p[1])-1][2]) + norm[2]
]
except:
self.vnorms[int(p[2])-1] = [norm[0],norm[1],norm[2]]
# Average out each vnorm..
index = 0
for v in self.vnorms.values():
self.vavenorms.append([v[0]/4, v[1]/4, v[2]/4])
index += 1
# Utility function to normalize a given vector
def normalize(self,v):
mag = v[0]*v[0]+v[1]*v[1]+v[2]*v[2]
if mag <= 0:
mag = 1
else:
mag = math.sqrt(mag)
return [v[0]/mag, v[1]/mag,v[2]/mag]
# Utility cross product function
def cross(self,v1,v2):
x = 0
y = 1
z = 2
return [v1[y]*v2[z] - v1[z]*v2[y],\
v1[z]*v2[x] - v1[x]*v2[z],\
v1[x]*v2[y] - v1[y]*v2[x]
]
# Constructor
def __init__(self, file,translate=[0,0,0],flip=0):
self.verts = []
self.polys = []
self.pnorms = []
self.vnorms = {}
self.vavenorms = []
self.translate = translate
self.flip = flip
print "Read in file: "+str(file)
self.readOBJ(file)
print "Done reading in WFMesh, now compute norms"
self.computeNorms()
print "Done computing norms, now display WFMesh"
# Draw Function
def __call__(self):
glColorMaterial(GL_FRONT, GL_DIFFUSE);
glEnable(GL_COLOR_MATERIAL);
glShadeModel(GL_SMOOTH);
# Color Everything grey
glColor3f(0.5,0.5,0.5);
glPushMatrix()
glTranslated(self.translate[0],self.translate[1],self.translate[2])
for index, p in enumerate(self.polys):
glBegin(GL_POLYGON)
glNormal3f(float(self.pnorms[index][0]),float(self.pnorms[index][1]),float(self.pnorms[index][2]))
for x in p:
glVertex3f(float(self.verts[int(x)-1][0]),float(self.verts[int(x)-1][1]),float(self.verts[int(x)-1][2]))
# Vertex Normals - not computed correctly, so commented out for now
# norm = self.vnorms[int(x)-1]
# glNormal3f(float(norm[0]),float(norm[1]),float(norm[2]))
glEnd()
glPopMatrix()
cmd.extend("createWFObj", createWFObj)