def remove_particles(infile, outfile, keep_rate):
"""Remove a percentage of particles from a .bgeo file.
Args:
infile: Path to input .bgeo file.
outfile: The file where the reduced number of particles will be stored.
keep_rate: A float between 0 and 1 that describes how much percentage of
particles will be kept in the new file.
"""
particles = partio.read(infile)
orig_attr = particles.attributeInfo("position")
orig_num_p = particles.numParticles()
all_indices = np.arange(orig_num_p)
keep_indices = np.random.choice(all_indices,
size=int(orig_num_p * keep_rate),
replace=False)
new_particles = partio.create()
P = new_particles.addAttribute("position", partio.VECTOR, 3)
id = new_particles.addAttribute("id", partio.INT, 1)
new_particles.addParticles(len(keep_indices))
for index, i in enumerate(keep_indices):
pos = particles.get(orig_attr, i)
new_particles.set(P, index, pos)
partio.write(outfile, new_particles)
def numpy_from_bgeo(path):
import partio
p = partio.read(path)
pos = p.attributeInfo('position')
vel = p.attributeInfo('velocity')
ida = p.attributeInfo('trackid') # old format
if ida is None:
ida = p.attributeInfo('id') # new format after splishsplash update
n = p.numParticles()
pos_arr = np.empty((n, pos.count))
for i in range(n):
pos_arr[i] = p.get(pos, i)
vel_arr = None
if not vel is None:
vel_arr = np.empty((n, vel.count))
for i in range(n):
vel_arr[i] = p.get(vel, i)
if not ida is None:
id_arr = np.empty((n, ), dtype=np.int64)
for i in range(n):
id_arr[i] = p.get(ida, i)[0]
s = np.argsort(id_arr)
result = [pos_arr[s]]
if not vel is None:
result.append(vel_arr[s])
else:
result = [pos_arr, vel_arr]
return tuple(result)
def setFile(self, filename):
""" Redraws the spreadsheet with a new particle file """
if filename == self.filename:
return
self.filename = filename
# reusing empty list causes column resizing problem so delete it
if self.treeView and not self.hasParticles:
self.layout().removeWidget(self.treeView)
del self.treeView
self.treeView = None
if not self.treeView:
self.treeView = QtWidgets.QTreeView()
self.treeView.setRootIsDecorated(False)
self.treeView.setUniformRowHeights(True)
self.layout().addWidget(self.treeView)
p = partio.read(filename)
if not p:
p = partio.create()
self.hasParticles = False
else:
self.hasParticles = True
model = PartioTreeModel(self.treeView, p, self.attrsOnly)
proxyModel = FilterModel(self.lineEdit, self.attrsOnly, self)
self.lineEdit.textChanged.connect(proxyModel.filter)
proxyModel.setSourceModel(model)
self.treeView.setModel(proxyModel)
for i in range(model.columnCount(self.treeView)):
self.treeView.resizeColumnToContents(i)
self.setWindowTitle(filename)
self.show()
def setFile(self, filename):
""" Redraws the spreadsheet with a new particle file """
if filename == self.filename:
return
self.filename = filename
# reusing empty list causes column resizing problem so delete it
if self.treeView and not self.hasParticles:
self.layout().removeWidget(self.treeView)
del self.treeView
self.treeView = None
if not self.treeView:
self.treeView = QtWidgets.QTreeView()
self.treeView.setRootIsDecorated(False)
self.treeView.setUniformRowHeights(True)
self.layout().addWidget(self.treeView)
p = partio.read(filename)
if not p:
p = partio.create()
self.hasParticles = False
else:
self.hasParticles = True
model = PartioTreeModel(self.treeView, p, self.attrsOnly)
proxyModel = FilterModel(self.lineEdit, self.attrsOnly, self)
self.lineEdit.textChanged.connect(proxyModel.filter)
proxyModel.setSourceModel(model)
self.treeView.setModel(proxyModel)
for i in range(model.columnCount(self.treeView)):
self.treeView.resizeColumnToContents(i)
self.setWindowTitle(filename)
self.show()
def testPartJson(self):
""" Test round-tripping """
testdir = os.path.dirname(os.path.abspath(__file__))
srcdir = os.path.dirname(testdir)
filename = os.path.join(srcdir, 'data', 'json.bgeo')
particleSet = partio.read(filename)
json1 = partjson.toJson(particleSet)
particleSet2 = partjson.fromJson(json1)
json2 = partjson.toJson(particleSet2)
self.assertEquals(json1, json2)
def read_partio(self, filename):
if filename[-7:] == '.pdb.gz':
global tempcache
import tempfile
import gzip
if filename in tempcache.keys():
ptc = partio.read(tempcache[filename].name)
else:
tmpf = tempfile.NamedTemporaryFile(suffix='.pdb32')
gzf = gzip.open(filename)
tmpf.write(gzf.read())
gzf.close()
ptc = partio.read(tmpf.name)
tempcache[filename] = tmpf
# tmpf.close() # python will close and remove temp files on exit
else:
ptc = partio.read(filename)
return ptc
def read_partio(self, filename):
if filename[-7:] == '.pdb.gz':
global tempcache
import tempfile
import gzip
if filename in tempcache.keys():
ptc = partio.read(tempcache[filename].name)
else:
tmpf = tempfile.NamedTemporaryFile(suffix='.pdb32')
gzf = gzip.open(filename)
tmpf.write(gzf.read())
gzf.close()
ptc = partio.read(tmpf.name)
tempcache[filename] = tmpf
# tmpf.close() # python will close and remove temp files on exit
else:
ptc = partio.read(filename)
return ptc
def partio_uncompress(dirname):
"""Take a folder of compressed .bgeo files and uncompress those files.
This can come in handy because SPlisHSPlasH stores it's particle data
as compressed .bgeo files.
Args:
dirname: Path to a directory of .bgeo files.
"""
for f in os.listdir(dirname):
if not f.endswith(".bgeo"):
continue
p = partio.read(os.path.join(dirname, f))
partio.write(os.path.join(dirname, f), p)
def read(self, filename):
""" Opens a file from disk and populates the UI """
if not os.path.exists(filename):
sys.stderr.write('Invalid filename: {}\n'.format(filename))
return
data = partio.read(filename)
if not data:
sys.stderr.write('Invalid particle file: {}\n'.format(filename))
data = partio.create()
self.filename = filename
self.setData(data)
self.setDirty(False)
def read(self, filename):
""" Opens a file from disk and populates the UI """
if not os.path.exists(filename):
sys.stderr.write('Invalid filename: {}\n'.format(filename))
return
data = partio.read(filename)
if not data:
sys.stderr.write('Invalid particle file: {}\n'.format(filename))
data = partio.create()
self.filename = filename
self.setData(data)
self.setDirty(False)
def numpy_from_bgeo(path):
import partio
p = partio.read(path)
for i in range(p.numAttributes()):
attr = p.attributeInfo(i)
typeStr = "NONE"
if attr.type == partio.VECTOR: typeStr = "VECTOR"
if attr.type == partio.FLOAT: typeStr = "FLOAT"
if attr.type == partio.INT: typeStr = "INT"
print("%10s[%d] %-30s " % (typeStr, attr.count, attr.name))
pos = p.attributeInfo('position')
vel = p.attributeInfo('velocity')
ida = p.attributeInfo('trackid') # old format
timestep = p.attributeInfo('timestep')
if ida is None:
ida = p.attributeInfo('id') # new format after splishsplash update
n = p.numParticles()
pos_arr = np.empty((n, pos.count))
for i in range(n):
pos_arr[i] = p.get(pos, i)
vel_arr = None
if not vel is None:
vel_arr = np.empty((n, vel.count))
for i in range(n):
vel_arr[i] = p.get(vel, i)
if not ida is None:
id_arr = np.empty((n, ), dtype=np.int64)
for i in range(n):
id_arr[i] = p.get(ida, i)[0]
s = np.argsort(id_arr)
result = [pos_arr[s]]
if not vel is None:
result.append(vel_arr[s])
result.append(timestep)
else:
result = [pos_arr, vel_arr, timestep]
return tuple(result)
def RequestData(self, request, inInfoVec, outInfoVec):
data_time = self.getCurrentTime(outInfoVec.GetInformationObject(0))
output = dsa.WrapDataObject(vtkUnstructuredGrid.GetData(outInfoVec))
currentFile = self.fileList[data_time]
p = partio.read(currentFile)
if p == None:
return 1
totalParticles = p.numParticles()
for i in range(p.numAttributes()):
attr = p.attributeInfo(i)
if attr.name == "position":
pos = np.array(p.data_buffer(attr), copy=False)
output.SetPoints(pos)
# cell_conn contains tuples (num indices, vertex ids)
# e.g. particles (1,0), (1,1), (1,2), (1,3), ...
# e.g. triangles (3, 0, 1, 2), (3, 2, 3, 4), ...
cell_conn = np.hstack([
np.ones((totalParticles, 1)),
np.arange(0, totalParticles, 1, dtype=int).reshape(-1, 1)
]).flatten()
# for particles use type VERTEX=1
cell_types = np.full((totalParticles), 1, np.ubyte)
# offset between two particles is 2 since cell_conn always contains the number of indices
cell_offsets = 2 * np.arange(totalParticles, dtype=int)
output.SetCells(cell_types, cell_offsets, cell_conn)
# add field data
for i in range(p.numAttributes()):
attr = p.attributeInfo(i)
if attr.name != "position":
values = np.array(p.data_buffer(attr), copy=False)
output.PointData.append(values, attr.name)
return 1
def main():
""" Main """
# Process command-line arguments
filenames = []
verbose = False
compress = False
for arg in sys.argv[1:]:
if arg in ('-h', '--help'):
print __doc__
return
if arg in ('-v', '--verbose'):
verbose = True
continue
if arg in ('-c', '--compress'):
compress = True
continue
filenames.append(arg)
if len(filenames) != 2:
print __doc__
sys.stderr.write('Incorrect number of arguments.\n')
sys.exit(1)
file1, file2 = filenames[0:2]
ext1 = os.path.splitext(file1)[1]
ext2 = os.path.splitext(file2)[1]
partio_extensions = ('.bgeo', '.geo', '.bhclassic', '.ptc', '.pdb')
# Validate files
if not os.path.exists(file1):
sys.stderr.write('Invalid input file: {}\n'.format(file1))
sys.exit(1)
# Convert from json to partio
if ext1 == '.json':
if ext2 not in partio_extensions:
sys.stderr.write('Unknown partio extension for: {}\n'.format(file2))
sys.exit(1)
with open(file1, 'r') as fp:
data = json.load(fp)
particleSet = fromJson(data)
partio.write(file2, particleSet, compress)
sys.exit(0)
if ext1 not in partio_extensions:
sys.stderr.write('Unknown partio extension for: {}\n'.format(file1))
sys.exit(1)
# Convert from partio to json
if ext1 in partio_extensions:
particleSet = partio.read(file1, verbose)
data = toJson(particleSet)
with open(file2, 'w') as fp:
json.dump(data, fp, indent=2, sort_keys=True)
sys.exit(0)
print __doc__
sys.stderr.write('Unknown file extension(s)')
sys.exit(1)
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND BASED ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
import sys
import partio
if __name__=="__main__":
filename=None
try:
filename=sys.argv[1]
except:
print "Usage: listAttr.py <filename>"
sys.exit(1)
# read particle name
p=partio.read(filename)
for i in range(p.numAttributes()):
attr=p.attributeInfo(i)
typeStr="NONE"
if attr.type==partio.VECTOR: typeStr="VECTOR"
if attr.type==partio.FLOAT: typeStr="FLOAT"
if attr.type==partio.INT: typeStr="INT"
print "%10s[%d] %-30s "%(typeStr,attr.count,attr.name)
def run(config):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # so the IDs match nvidia-smi
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu_id # "0, 1" for multiple
prepare_dirs_and_logger(config)
tf.compat.v1.set_random_seed(config.seed)
config.rng = np.random.RandomState(config.seed)
styler = Styler(config)
styler.load_img(config.resolution)
params = {}
# load particles
p = []
r = []
for i in trange(config.num_frames, desc='load particle'):
pt_path = os.path.join(config.data_dir, config.dataset, config.d_path % (config.target_frame+i))
pt = partio.read(pt_path)
p_id = pt.attributeInfo('id')
p_pos = pt.attributeInfo('position')
p_den = pt.attributeInfo('density')
p_num = pt.numParticles()
p_ = np.zeros([p_num,2], dtype=np.float32)
r_ = np.zeros([p_num,1], dtype=np.float32)
for j in range(p_num):
p_id_ = pt.get(p_id, j)[0]
p_[p_id_] = pt.get(p_pos, p_id_)[:-1] # 2d
r_[p_id_] = pt.get(p_den, p_id_)
r.append(r_)
# normalize particle position [0-1]
px, py = p_[...,0], p_[...,1]
px /= config.domain[1]
py /= config.domain[0]
p_ = np.stack([py,px], axis=-1)
p.append(p_)
print('resolution:', config.resolution)
print('domain:', config.domain)
print('radius:', config.radius)
print('normalized px range', px.min(), px.max())
print('normalized py range', py.min(), py.max())
print('num particles:', p[0].shape) # the number of particles is fixed
params['p'] = p
params['r'] = r
# styler.render_test(params)
result = styler.run(params)
# save loss plot
l = result['l']
lb = []
for o, l_ in enumerate(l):
lb_, = plt.plot(range(len(l_)), l_, label='oct %d' % o)
lb.append(lb_)
plt.legend(handles=lb)
# plt.show()
plot_path = os.path.join(config.log_dir, 'loss_plot.png')
plt.savefig(plot_path)
# save density fields
d_sty = result['d'] # [0-255], uint8
# d_path = os.path.join(config.log_dir, 'd%03d_%03d.png' % (config.target_frame,config.target_frame+config.num_frames-1))
# save_image(d_sty, d_path, nrow=5, gray=not 'c' in config.target_field)
for i, d_sty_ in enumerate(d_sty):
im = Image.fromarray(d_sty_)
d_path = os.path.join(config.log_dir, '%03d.png' % (config.target_frame+i))
im.save(d_path)
d_intm = result['d_intm']
for o, d_intm_o in enumerate(d_intm):
for i, d_intm_ in enumerate(d_intm_o):
im = Image.fromarray(d_intm_)
d_path = os.path.join(config.log_dir, 'o%02d_%03d.png' % (o, config.target_frame))
im.save(d_path)
# save particles (load using Houdini GPlay)
c_sty = result['c']
p_org = []
for p_ in p:
# denormalize particle positions
px, py = p_[...,1], p_[...,0]
px *= config.domain[1]
py *= config.domain[0]
p_org.append(np.stack([px,py], axis=-1))
for i in range(config.num_frames):
# create a particle set and attributes
pt = partio.create()
position = pt.addAttribute("position",partio.VECTOR,2)
color = pt.addAttribute("Cd",partio.FLOAT,3)
radius = pt.addAttribute("radius",partio.FLOAT,1)
# normal = pt.addAttribute("normal",partio.VECTOR,3)
for pi in range(p_org[i].shape[0]):
p_ = pt.addParticle()
pt.set(position, p_, tuple(p_org[i][pi].astype(np.float)))
pt.set(color, p_, tuple(c_sty[i][pi].astype(np.float)))
pt.set(radius, p_, (config.radius,))
p_path = os.path.join(config.log_dir, '%03d.bgeo' % (config.target_frame+i))
partio.write(p_path, pt)
# visualization using open3d
bbox = [
[0,0,-1],
[config.domain[1],config.domain[0],1], # [X,Y,Z]
]
draw_pt(p_org, pc=c_sty, bbox=bbox, dt=0.1)
def run(config):
os.environ[
"CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # so the IDs match nvidia-smi
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu_id # "0, 1" for multiple
prepare_dirs_and_logger(config)
tf.compat.v1.set_random_seed(config.seed)
config.rng = np.random.RandomState(config.seed)
styler = Styler(config)
styler.load_img(config.resolution[1:])
params = {}
# load particles
nmin, nmax = np.iinfo(np.int32).max, 0
for i in range(config.num_frames):
pt_path = os.path.join(config.data_dir, config.dataset,
config.d_path % (config.target_frame + i))
pt = partio.read(pt_path)
p_num = pt.numParticles()
nmin = min(p_num, nmin)
nmax = max(p_num, nmax)
print('# range:', nmin, nmax)
p = []
# r = []
for i in trange(config.num_frames,
desc='load particle'): # last one for mask
pt_path = os.path.join(config.data_dir, config.dataset,
config.d_path % (config.target_frame + i))
pt = partio.read(pt_path)
p_attr_id = pt.attributeInfo('id')
p_attr_pos = pt.attributeInfo('position')
# p_attr_den = pt.attributeInfo('density')
p_ = np.ones([nmax, 3], dtype=np.float32) * -1
# r_ = np.zeros([nmax,1], dtype=np.float32)
p_num = pt.numParticles()
for j in range(p_num):
p_id_j = pt.get(p_attr_id, j)[0]
p_[p_id_j] = pt.get(p_attr_pos, p_id_j)
# r_[p_id_j] = pt.get(p_attr_den, p_id_j)
# r.append(r_)
# normalize particle position [0-1]
px, py, pz = p_[..., 0], p_[..., 1], p_[..., 2]
px /= config.domain[2]
py /= config.domain[1]
pz /= config.domain[0]
p_ = np.stack([pz, py, px], axis=-1)
p.append(p_)
print('resolution:', config.resolution)
print('domain:', config.domain)
print('radius:', config.radius)
print('normalized px range', px.min(), px.max())
print('normalized py range', py.min(), py.max())
params['p'] = p
# styler.render_test(params)
result = styler.run(params)
# save loss plot
l = result['l']
lb = []
for o, l_ in enumerate(l):
lb_, = plt.plot(range(len(l_)), l_, label='oct %d' % o)
lb.append(lb_)
plt.legend(handles=lb)
# plt.show()
plot_path = os.path.join(config.log_dir, 'loss_plot.png')
plt.savefig(plot_path)
# save particle (load using Houdini GPlay)
p_sty = result['p']
p = []
# v_sty = result['v']
# v = []
for i in range(config.num_frames):
# denormalize particle positions
px, py, pz = p_sty[i][..., 2], p_sty[i][..., 1], p_sty[i][..., 0]
px *= config.domain[2]
py *= config.domain[1]
pz *= config.domain[0]
p_sty_ = np.stack([px, py, pz], axis=-1)
p.append(p_sty_)
# # denormalize particle displacement for stylization
# vx, vy, vz = v_sty[i][...,2], v_sty[i][...,1], v_sty[i][...,0]
# vx *= config.domain[2]
# vy *= config.domain[1]
# vz *= config.domain[0]
# v_sty_ = np.stack([vx,vy,vz], axis=-1)
# v.append(v_sty_)
# create a particle set and attributes
pt = partio.create()
position = pt.addAttribute("position", partio.VECTOR, 3)
# color = pt.addAttribute("Cd",partio.FLOAT,3)
radius = pt.addAttribute("radius", partio.FLOAT, 1)
# normal = pt.addAttribute("normal",partio.VECTOR,3)
for p_sty_i in p_sty_:
if p_sty_i[0] < 0: continue
p_ = pt.addParticle()
pt.set(position, p_, tuple(p_sty_i.astype(np.float)))
pt.set(radius, p_, (config.radius, ))
p_path = os.path.join(config.log_dir,
'%03d.bgeo' % (config.target_frame + i))
partio.write(p_path, pt)
r_sty = result['r']
for i, r_sty_ in enumerate(r_sty):
im = Image.fromarray(r_sty_)
d_path = os.path.join(config.log_dir,
'%03d.png' % (config.target_frame + i))
im.save(d_path)
d_intm = result['d_intm']
for o, d_intm_o in enumerate(d_intm):
for i, d_intm_ in enumerate(d_intm_o):
if d_intm_ is None: continue
im = Image.fromarray(d_intm_)
d_path = os.path.join(
config.log_dir,
'o%02d_%03d.png' % (o, config.target_frame + i))
im.save(d_path)
# visualization using open3d
bbox = [
[0, 0, 0],
[config.domain[2], config.domain[1], config.domain[0]], # [X,Y,Z]
]
draw_pt(p, bbox=bbox, dt=0.1, is_2d=False) # pv=v,
import partio
import os
fileLocation = os.path.dirname(__file__)
print fileLocation
for i in range (0,100):
p=partio.read(fileLocation+"SnowF"+str(i).zfill(4)+".ptc")
partio.write(fileLocation+"SnowF"+str(i).zfill(4)+".bgeo",p)
def main():
""" Main """
# Process command-line arguments
filenames = []
verbose = False
compress = False
for arg in sys.argv[1:]:
if arg in ('-h', '--help'):
print(__doc__)
return
if arg in ('-v', '--verbose'):
verbose = True
continue
if arg in ('-c', '--compress'):
compress = True
continue
filenames.append(arg)
if len(filenames) != 2:
print(__doc__)
sys.stderr.write('Incorrect number of arguments.\n')
sys.exit(1)
file1, file2 = filenames[0:2]
ext1 = os.path.splitext(file1)[1]
ext2 = os.path.splitext(file2)[1]
partio_extensions = ('.bgeo', '.geo', '.bhclassic', '.ptc', '.pdb')
# Validate files
if not os.path.exists(file1):
sys.stderr.write('Invalid input file: {}\n'.format(file1))
sys.exit(1)
# Convert from json to partio
if ext1 == '.json':
if ext2 not in partio_extensions:
sys.stderr.write(
'Unknown partio extension for: {}\n'.format(file2))
sys.exit(1)
with open(file1, 'r') as fp:
data = json.load(fp)
particleSet = fromJson(data)
partio.write(file2, particleSet, compress)
sys.exit(0)
if ext1 not in partio_extensions:
sys.stderr.write('Unknown partio extension for: {}\n'.format(file1))
sys.exit(1)
# Convert from partio to json
if ext1 in partio_extensions:
particleSet = partio.read(file1, verbose)
data = toJson(particleSet)
with open(file2, 'w') as fp:
json.dump(data, fp, indent=2, sort_keys=True)
sys.exit(0)
print(__doc__)
sys.stderr.write('Unknown file extension(s)')
sys.exit(1)
def __call__(self, scene, depsgraph=None):
partioFile = self.param[0]
emitterObject = self.param[1]
try:
dummy = emitterObject.name
except:
# emitter does not exist anymore
#clear the post frame handler
bpy.app.handlers.frame_change_post.remove(self)
return
indexlist = re.findall(r'\d+', partioFile)
self.isSequence = True
if len(indexlist) == 0:
self.isSequence = False
fileName = partioFile
else:
frameNumber = int(indexlist[-1])
idx = partioFile.rfind(str(frameNumber))
l = len(str(frameNumber))
fileName = str(
partioFile[0:idx]) + str(scene.frame_current - 1) + str(
partioFile[idx + l:])
print("Read partio file: " + fileName)
p = partio.read(fileName)
if p != None:
totalParticles = p.numParticles()
print("# particles: " + str(totalParticles))
emitterObject.particle_systems[0].settings.count = totalParticles
if totalParticles > 10000:
emitterObject.particle_systems[
0].settings.display_method = 'DOT'
if depsgraph is None:
depsgraph = bpy.context.evaluated_depsgraph_get()
particle_systems = emitterObject.evaluated_get(
depsgraph).particle_systems
particles = particle_systems[0].particles
posAttr = None
velAttr = None
for i in range(p.numAttributes()):
attr = p.attributeInfo(i)
typeStr = "NONE"
if attr.name == "position": posAttr = attr
if attr.name == "velocity": velAttr = attr
pos = []
for i in range(p.numParticles()):
position = p.get(posAttr, i)
x = mathutils.Vector((position[0], -position[2], position[1]))
x = emitterObject.matrix_world @ x
pos.append(x[0])
pos.append(x[1])
pos.append(x[2])
# Set the location of all particle locations to flatList
particles.foreach_set("location", pos)
if velAttr is not None:
vel = []
for i in range(p.numParticles()):
velocity = p.get(velAttr, i)
v = mathutils.Vector(
(velocity[0], -velocity[2], velocity[1]))
v = emitterObject.matrix_world @ v - emitterObject.location
vel.append(v[0])
vel.append(v[1])
vel.append(v[2])
particles.foreach_set("velocity", vel)
emitterObject.particle_systems[0].settings.frame_end = 0
def count_particles(filename):
"""Count the number of particles in a .bgeo file.
"""
return partio.read(filename).numParticles()
def run(config):
os.environ[
"CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # so the IDs match nvidia-smi
os.environ["CUDA_VISIBLE_DEVICES"] = config.gpu_id # "0, 1" for multiple
prepare_dirs_and_logger(config)
tf.compat.v1.set_random_seed(config.seed)
config.rng = np.random.RandomState(config.seed)
styler = Styler(config)
styler.load_img(config.resolution[1:])
params = {}
# the number of particles range
nmin, nmax = np.iinfo(np.int32).max, 0
for i in range(config.num_frames):
pt_path = os.path.join(config.data_dir, config.dataset,
config.d_path % (config.target_frame + i))
pt = partio.read(pt_path)
p_num = pt.numParticles()
nmin, nmax = min(nmin, p_num), max(nmax, p_num)
print('# range:', nmin, nmax)
p, r = [], []
for i in trange(config.num_frames, desc='load particle'):
pt_path = os.path.join(config.data_dir, config.dataset,
config.d_path % (config.target_frame + i))
pt = partio.read(pt_path)
p_id = pt.attributeInfo('id')
p_pos = pt.attributeInfo('position')
p_den = pt.attributeInfo('density')
p_ = np.ones([nmax, 3], dtype=np.float32) * -1
r_ = np.zeros([nmax, config.num_kernels], dtype=np.float32)
p_num = pt.numParticles()
for j in range(p_num):
p_id_j = pt.get(p_id, j)[0]
p_[j] = pt.get(p_pos, p_id_j)
r_[j] = pt.get(p_den, p_id_j)
r.append(r_)
# normalize particle position [0-1]
px, py, pz = p_[..., 0], p_[..., 1], p_[..., 2]
px /= config.domain[2]
py /= config.domain[1]
pz /= config.domain[0]
p_ = np.stack([pz, py, px], axis=-1)
p.append(p_)
print('resolution:', config.resolution)
print('domain:', config.domain)
print('radius:', config.radius)
print('normalized px range', px.min(), px.max())
print('normalized py range', py.min(), py.max())
print('normalized pz range', pz.min(), pz.max())
params['p'] = p
params['r'] = r
# styler.render_test(params)
result = styler.run(params)
# save loss plot
l = result['l']
lb = []
for o, l_ in enumerate(l):
lb_, = plt.plot(range(len(l_)), l_, label='oct %d' % o)
lb.append(lb_)
plt.legend(handles=lb)
# plt.show()
plot_path = os.path.join(config.log_dir, 'loss_plot.png')
plt.savefig(plot_path)
r_sty = result['r']
for i, r_sty_ in enumerate(r_sty):
im = Image.fromarray(r_sty_)
d_path = os.path.join(config.log_dir,
'%03d.png' % (config.target_frame + i))
im.save(d_path)
d_sty = result['d']
for i, d_sty_ in enumerate(d_sty):
d_path = os.path.join(config.log_dir,
'%03d.npz' % (config.target_frame + i))
np.savez_compressed(d_path, x=d_sty_[:, ::-1])
d_intm = result['d_intm']
for o, d_intm_o in enumerate(d_intm):
for i, d_intm_ in enumerate(d_intm_o):
if d_intm_ is None: continue
im = Image.fromarray(d_intm_)
d_path = os.path.join(
config.log_dir,
'o%02d_%03d.png' % (o, config.target_frame + i))
im.save(d_path)
