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 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 partio_write_rigid_body(vertices, filename):
"""Write vertices of an object into a .bgeo file.
Args:
vertices: A num_vertices x 3 numpy array.
filename: Path to the .bgeo file to create.
"""
os.makedirs(os.path.dirname(filename), exist_ok=True)
particleSet = partio.create()
P = particleSet.addAttribute("position", partio.VECTOR, 3)
id = particleSet.addAttribute("id", partio.INT, 1)
particleSet.addParticles(len(vertices))
for i, vertex in enumerate(vertices):
particleSet.set(P, i, [float(x) for x in vertex])
partio.write(filename, particleSet)
def write_bgeo_from_numpy(outpath, pos_arr, vel_arr):
import partio
n = pos_arr.shape[0]
if not (vel_arr.shape[0] == n and pos_arr.shape[1] == 3
and vel_arr.shape[1] == 3):
raise ValueError(
"invalid shapes for pos_arr {} and/or vel_arr {}".format(
pos_arr.shape, vel_arr.shape))
p = partio.create()
position_attr = p.addAttribute("position", partio.VECTOR, 3)
velocity_attr = p.addAttribute("velocity", partio.VECTOR, 3)
for i in range(n):
idx = p.addParticle()
p.set(position_attr, idx, pos_arr[i].astype(float))
p.set(velocity_attr, idx, vel_arr[i].astype(float))
partio.write(outpath, p)
def write(self, filename, delta):
""" Write data to file. If delta is False, saves a full copy
of the data, rebaselining. If delta is True, saves only
the particles (todo: and attributes) that have changed,
but maintains the original baseline
"""
if not self.data:
return
# If we're saving a delta, create a new particle set with just
# the differences from the original.
if delta:
data = self.createDelta()
else:
data = self.data
partio.write(filename, data)
# If we saved a full copy, rebaseline
if not delta:
self.filename = filename
self.originalData = copy(data)
self.setDirty(False)
def write(self, filename, delta):
""" Write data to file. If delta is False, saves a full copy
of the data, rebaselining. If delta is True, saves only
the particles (todo: and attributes) that have changed,
but maintains the original baseline
"""
if not self.data:
return
# If we're saving a delta, create a new particle set with just
# the differences from the original.
if delta:
data = self.createDelta()
else:
data = self.data
partio.write(filename, data)
# If we saved a full copy, rebaseline
if not delta:
self.filename = filename
self.originalData = copy(data)
self.setDirty(False)
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 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,
for pobj in particles:
pset = partio.create()
pattr = pset.addAttribute("position", partio.VECTOR, 3)
cattr = pset.addAttribute("color", partio.VECTOR, 3)
nattr = pset.addAttribute("normal", partio.VECTOR, 3)
pcount = cmds.getAttr(pobj + '.count')
pset.addParticles(pcount)
print('[Particle Export] Exporting %s (%d points)...' % (pobj, pcount))
for i in range(pcount):
# get position
pos = cmds.getParticleAttr(pobj + '.pt[%d]' % i, at='position')[0:3]
# get color
if cmds.attributeQuery('rgbPP', n=pobj, ex=True):
col = cmds.getParticleAttr(pobj + '.pt[%d]' % i, at='rgbPP')[0:3]
elif cmds.attributeQuery('colorRed', n=pobj, ex=True):
col = (cmds.getAttr(pobj + '.colorRed'),
cmds.getAttr(pobj + '.colorGreen'),
cmds.getAttr(pobj + '.colorBlue'))
else:
col = (1, 0, 0)
# normal
nml = (0, 0, 0)
# set attributes
pset.set(pattr, i, pos)
pset.set(cattr, i, col)
pset.set(nattr, i, nml)
partio.write(str('%s/%s.bgeo' % (path, pobj)), pset, True)
print('[Particle Export] Written to %s/%s.bgeo' % (path, pobj))
print('-' * 50)
print('[Particle Export] Done %d objects.' % len(particles))
cmds.select(selected)
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)
# create a particle set and attributes
p=partio.create()
position=p.addAttribute("position",partio.VECTOR,3)
color=p.addAttribute("Cd",partio.FLOAT,3)
radius=p.addAttribute("radius",partio.FLOAT,1)
normal=p.addAttribute("normal",partio.VECTOR,3)
# make a spiral circle thingy
t=0
dt=.01
import math
while t<2*math.pi:
# allocate new particle
i=p.addParticle()
# set particle attributes
r=.25*math.sin(10*t)+1.
x,y,z=-r*math.sin(t),0,r*math.cos(t)
p.set(position,i,(x,y,z))
p.set(color,i,(t/2./math.pi,1-(t/2./math.pi),0))
p.set(radius,i,(.02,))
p.set(normal,i,(0,1,0))
foo=p.get(position,i)
t+=dt
# Write to a point cloud
partio.write("spiral.ptc",p)
# create a particle set and attributes
p = partio.create()
position = p.addAttribute("position", partio.VECTOR, 3)
color = p.addAttribute("Cd", partio.FLOAT, 3)
radius = p.addAttribute("radius", partio.FLOAT, 1)
normal = p.addAttribute("normal", partio.VECTOR, 3)
# make a spiral circle thingy
t = 0
dt = .01
import math
while t < 2 * math.pi:
# allocate new particle
i = p.addParticle()
# set particle attributes
r = .25 * math.sin(10 * t) + 1.
x, y, z = -r * math.sin(t), 0, r * math.cos(t)
p.set(position, i, (x, y, z))
p.set(color, i, (t / 2. / math.pi, 1 - (t / 2. / math.pi), 0))
p.set(radius, i, (.02, ))
p.set(normal, i, (0, 1, 0))
foo = p.get(position, i)
t += dt
# Write to a point cloud
partio.write("spiral.ptc", p)
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 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)
resampler = SimG2P(config)
# load input density fields
for t in trange(config.num_frames,
desc='load density'): # last one for mask
d_path = os.path.join(config.data_dir, config.dataset,
config.d_path % (config.target_frame + t))
with np.load(d_path) as data:
d = data['x'][:, ::-1] # [D,H,W], [0-1]
# mantaflow dataset
v_path = os.path.join(config.data_dir, config.dataset,
config.v_path % (config.target_frame + t))
with np.load(v_path) as data:
v_ = data['x'] # [D,H,W,3]
vx = np.dstack(
(v_, np.zeros((v_.shape[0], v_.shape[1], 1, v_.shape[3]))))
vx = (vx[:, :, 1:, 0] + vx[:, :, :-1, 0]) * 0.5
vy = np.hstack(
(v_, np.zeros((v_.shape[0], 1, v_.shape[2], v_.shape[3]))))
vy = (vy[:, 1:, :, 1] + vy[:, :-1, :, 1]) * 0.5
vz = np.vstack(
(v_, np.zeros((1, v_.shape[1], v_.shape[2], v_.shape[3]))))
vz = (vz[1:, :, :, 2] + vz[:-1, :, :, 2]) * 0.5
v_ = np.stack([vx, vy, vz], axis=-1)
v_ = v_[:, ::-1]
vx = v_[..., 0] / v_.shape[2] * config.scale
vy = -v_[..., 1] / v_.shape[1] * config.scale
vz = v_[..., 2] / v_.shape[0] * config.scale
u = np.stack([vz, vy, vx], axis=-1)
if config.resampling:
if t == 0:
n_prev = 0
# sampling at the beginning wo opt.
p, p_id = resampler.sample(d, disc=config.disc, threshold=0)
result = resampler.optimize(p, p_id, d, u)
p = result['p']
p_id = result['p_id']
p_den = result['p_den']
# d_diff = result['d_diff']
# plt.imshow(d_diff); plt.show()
l = result['l'][-1] # last loss
d_smp = result['d_smp']
else:
if t == 0:
n_prev = 0
# sampling at the beginning wo opt.
p, p_id = resampler.sample(d, disc=config.disc, threshold=0)
p_src = p
else:
# simply source particles of t=0
p = np.concatenate([p, p_src], axis=0)
p_id = np.arange(p.shape[0])
p_den = np.ones([p.shape[0], 1])
p, d_smp = resampler.naive_adv(p, u, p_den)
l = 0
print(t, 'num particles', p.shape[0], '(+%d)' % (p.shape[0] - n_prev),
'loss', l)
n_prev = p.shape[0]
# convert to the original domain coordinate
px, py, pz = p[..., 2], 1 - p[..., 1], p[..., 0]
p_ = np.stack([
px * config.domain[2], py * config.domain[1], pz * config.domain[0]
],
axis=-1)
# create a particle set and attributes
pt = partio.create()
pid = pt.addAttribute('id', partio.INT, 1)
position = pt.addAttribute("position", partio.VECTOR, 3)
if p_den.shape[1] > 1:
density = pt.addAttribute('density', partio.VECTOR, p_den.shape[1])
else:
density = pt.addAttribute('density', partio.FLOAT, 1)
color = pt.addAttribute("Cd", partio.FLOAT, 3)
radius = pt.addAttribute("radius", partio.FLOAT, 1)
for i in range(p_.shape[0]):
pt_ = pt.addParticle()
pt.set(pid, pt_, (int(p_id[i]), ))
pt.set(position, pt_, tuple(p_[i].astype(np.float)))
if p_den.shape[1] > 1:
pt.set(density, pt_, tuple(p_den[i].astype(np.float)))
else:
pt.set(density, pt_, (float(p_den[i]), ))
pt.set(color, pt_,
tuple(np.array([p_den[i, 0]] * 3, dtype=np.float)))
pt.set(radius, pt_, (config.radius, ))
# save particle
p_path = os.path.join(config.log_dir,
'%03d.bgeo' % (config.target_frame + t))
partio.write(p_path, pt)
# save density image
transmit = np.exp(-np.cumsum(d_smp[::-1], axis=0) * config.transmit)
d_img = np.sum(d_smp * transmit, axis=0)
d_img /= d_img.max()
im = Image.fromarray((d_img[::-1] * 255).astype(np.uint8))
im_path = os.path.join(config.log_dir,
'%03d.png' % (config.target_frame + t))
im.save(im_path)
stat_path = os.path.join(config.log_dir, 'stat.txt')
with open(stat_path, 'w') as f:
f.write('num particles %d\n' % p.shape[0])
f.write('loss %.2f' % l)
