function_name,docstring,function_body,file_path
diffusion_from_config,,"def diffusion_from_config(config: Dict[str, Any]) ->GaussianDiffusion:
schedule = config['schedule']
steps = config['timesteps']
respace = config.get('respacing', None)
mean_type = config.get('mean_type', 'epsilon')
betas = get_named_beta_schedule(schedule, steps)
channel_scales = config.get('channel_scales', None)
channel_biases = config.get('channel_biases', None)
if channel_scales is not None:
channel_scales = np.array(channel_scales)
if channel_biases is not None:
channel_biases = np.array(channel_biases)
kwargs = dict(betas=betas, model_mean_type=mean_type, model_var_type=
'learned_range', loss_type='mse', channel_scales=channel_scales,
channel_biases=channel_biases)
if respace is None:
return GaussianDiffusion(**kwargs)
else:
return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace
), **kwargs)
",point_e\diffusion\configs.py
get_beta_schedule,"This is the deprecated API for creating beta schedules.
See get_named_beta_schedule() for the new library of schedules.","def get_beta_schedule(beta_schedule, *, beta_start, beta_end,
num_diffusion_timesteps):
""""""""""""
if beta_schedule == 'linear':
betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps,
dtype=np.float64)
else:
raise NotImplementedError(beta_schedule)
assert betas.shape == (num_diffusion_timesteps,)
return betas
",point_e\diffusion\gaussian_diffusion.py
get_named_beta_schedule,"Get a pre-defined beta schedule for the given name.
The beta schedule library consists of beta schedules which remain similar
in the limit of num_diffusion_timesteps.
Beta schedules may be added, but should not be removed or changed once
they are committed to maintain backwards compatibility.","def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
""""""""""""
if schedule_name == 'linear':
scale = 1000 / num_diffusion_timesteps
return get_beta_schedule('linear', beta_start=scale * 0.0001,
beta_end=scale * 0.02, num_diffusion_timesteps=
num_diffusion_timesteps)
elif schedule_name == 'cosine':
return betas_for_alpha_bar(num_diffusion_timesteps, lambda t: math.
cos((t + 0.008) / 1.008 * math.pi / 2) ** 2)
else:
raise NotImplementedError(f'unknown beta schedule: {schedule_name}')
",point_e\diffusion\gaussian_diffusion.py
betas_for_alpha_bar,"Create a beta schedule that discretizes the given alpha_t_bar function,
which defines the cumulative product of (1-beta) over time from t = [0,1].
:param num_diffusion_timesteps: the number of betas to produce.
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
produces the cumulative product of (1-beta) up to that
part of the diffusion process.
:param max_beta: the maximum beta to use; use values lower than 1 to
prevent singularities.","def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
""""""""""""
betas = []
for i in range(num_diffusion_timesteps):
t1 = i / num_diffusion_timesteps
t2 = (i + 1) / num_diffusion_timesteps
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
return np.array(betas)
",point_e\diffusion\gaussian_diffusion.py
space_timesteps,"Create a list of timesteps to use from an original diffusion process,
given the number of timesteps we want to take from equally-sized portions
of the original process.
For example, if there's 300 timesteps and the section counts are [10,15,20]
then the first 100 timesteps are strided to be 10 timesteps, the second 100
are strided to be 15 timesteps, and the final 100 are strided to be 20.
:param num_timesteps: the number of diffusion steps in the original
process to divide up.
:param section_counts: either a list of numbers, or a string containing
comma-separated numbers, indicating the step count
per section. As a special case, use ""ddimN"" where N
is a number of steps to use the striding from the
DDIM paper.
:return: a set of diffusion steps from the original process to use.","def space_timesteps(num_timesteps, section_counts):
""""""""""""
if isinstance(section_counts, str):
if section_counts.startswith('ddim'):
desired_count = int(section_counts[len('ddim'):])
for i in range(1, num_timesteps):
if len(range(0, num_timesteps, i)) == desired_count:
return set(range(0, num_timesteps, i))
raise ValueError(
f'cannot create exactly {num_timesteps} steps with an integer stride'
)
elif section_counts.startswith('exact'):
res = set(int(x) for x in section_counts[len('exact'):].split(','))
for x in res:
if x < 0 or x >= num_timesteps:
raise ValueError(f'timestep out of bounds: {x}')
return res
section_counts = [int(x) for x in section_counts.split(',')]
size_per = num_timesteps // len(section_counts)
extra = num_timesteps % len(section_counts)
start_idx = 0
all_steps = []
for i, section_count in enumerate(section_counts):
size = size_per + (1 if i < extra else 0)
if size < section_count:
raise ValueError(
f'cannot divide section of {size} steps into {section_count}')
if section_count <= 1:
frac_stride = 1
else:
frac_stride = (size - 1) / (section_count - 1)
cur_idx = 0.0
taken_steps = []
for _ in range(section_count):
taken_steps.append(start_idx + round(cur_idx))
cur_idx += frac_stride
all_steps += taken_steps
start_idx += size
return set(all_steps)
",point_e\diffusion\gaussian_diffusion.py
_extract_into_tensor,"Extract values from a 1-D numpy array for a batch of indices.
:param arr: the 1-D numpy array.
:param timesteps: a tensor of indices into the array to extract.
:param broadcast_shape: a larger shape of K dimensions with the batch
dimension equal to the length of timesteps.
:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.","def _extract_into_tensor(arr, timesteps, broadcast_shape):
""""""""""""
res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
while len(res.shape) < len(broadcast_shape):
res = res[..., None]
return res + th.zeros(broadcast_shape, device=timesteps.device)
",point_e\diffusion\gaussian_diffusion.py
normal_kl,"Compute the KL divergence between two gaussians.
Shapes are automatically broadcasted, so batches can be compared to
scalars, among other use cases.","def normal_kl(mean1, logvar1, mean2, logvar2):
""""""""""""
tensor = None
for obj in (mean1, logvar1, mean2, logvar2):
if isinstance(obj, th.Tensor):
tensor = obj
break
assert tensor is not None, 'at least one argument must be a Tensor'
logvar1, logvar2 = [(x if isinstance(x, th.Tensor) else th.tensor(x).to
(tensor)) for x in (logvar1, logvar2)]
return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + (
mean1 - mean2) ** 2 * th.exp(-logvar2))
",point_e\diffusion\gaussian_diffusion.py
approx_standard_normal_cdf,"A fast approximation of the cumulative distribution function of the
standard normal.","def approx_standard_normal_cdf(x):
""""""""""""
return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.
pow(x, 3))))
",point_e\diffusion\gaussian_diffusion.py
discretized_gaussian_log_likelihood,"Compute the log-likelihood of a Gaussian distribution discretizing to a
given image.
:param x: the target images. It is assumed that this was uint8 values,
rescaled to the range [-1, 1].
:param means: the Gaussian mean Tensor.
:param log_scales: the Gaussian log stddev Tensor.
:return: a tensor like x of log probabilities (in nats).","def discretized_gaussian_log_likelihood(x, *, means, log_scales):
""""""""""""
assert x.shape == means.shape == log_scales.shape
centered_x = x - means
inv_stdv = th.exp(-log_scales)
plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
cdf_plus = approx_standard_normal_cdf(plus_in)
min_in = inv_stdv * (centered_x - 1.0 / 255.0)
cdf_min = approx_standard_normal_cdf(min_in)
log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
cdf_delta = cdf_plus - cdf_min
log_probs = th.where(x < -0.999, log_cdf_plus, th.where(x > 0.999,
log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))))
assert log_probs.shape == x.shape
return log_probs
",point_e\diffusion\gaussian_diffusion.py
mean_flat,Take the mean over all non-batch dimensions.,"def mean_flat(tensor):
""""""""""""
return tensor.flatten(1).mean(1)
",point_e\diffusion\gaussian_diffusion.py
__init__,,"def __init__(self, *, betas: Sequence[float], model_mean_type: str,
model_var_type: str, loss_type: str, discretized_t0: bool=False,
channel_scales: Optional[np.ndarray]=None, channel_biases: Optional[np.
ndarray]=None):
self.model_mean_type = model_mean_type
self.model_var_type = model_var_type
self.loss_type = loss_type
self.discretized_t0 = discretized_t0
self.channel_scales = channel_scales
self.channel_biases = channel_biases
betas = np.array(betas, dtype=np.float64)
self.betas = betas
assert len(betas.shape) == 1, 'betas must be 1-D'
assert (betas > 0).all() and (betas <= 1).all()
self.num_timesteps = int(betas.shape[0])
alphas = 1.0 - betas
self.alphas_cumprod = np.cumprod(alphas, axis=0)
self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (
1.0 - self.alphas_cumprod)
self.posterior_log_variance_clipped = np.log(np.append(self.
posterior_variance[1], self.posterior_variance[1:]))
self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (
1.0 - self.alphas_cumprod)
self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(
alphas) / (1.0 - self.alphas_cumprod)
",point_e\diffusion\gaussian_diffusion.py
get_sigmas,,"def get_sigmas(self, t):
return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)
",point_e\diffusion\gaussian_diffusion.py
q_mean_variance,"Get the distribution q(x_t | x_0).
:param x_start: the [N x C x ...] tensor of noiseless inputs.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:return: A tuple (mean, variance, log_variance), all of x_start's shape.","def q_mean_variance(self, x_start, t):
""""""""""""
mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
) * x_start
variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape
)
log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod,
t, x_start.shape)
return mean, variance, log_variance
",point_e\diffusion\gaussian_diffusion.py
q_sample,"Diffuse the data for a given number of diffusion steps.
In other words, sample from q(x_t | x_0).
:param x_start: the initial data batch.
:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
:param noise: if specified, the split-out normal noise.
:return: A noisy version of x_start.","def q_sample(self, x_start, t, noise=None):
""""""""""""
if noise is None:
noise = th.randn_like(x_start)
assert noise.shape == x_start.shape
return _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
) * x_start + _extract_into_tensor(self.
sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
",point_e\diffusion\gaussian_diffusion.py
q_posterior_mean_variance,"Compute the mean and variance of the diffusion posterior:
q(x_{t-1} | x_t, x_0)","def q_posterior_mean_variance(self, x_start, x_t, t):
""""""""""""
assert x_start.shape == x_t.shape
posterior_mean = _extract_into_tensor(self.posterior_mean_coef1, t, x_t
.shape) * x_start + _extract_into_tensor(self.posterior_mean_coef2,
t, x_t.shape) * x_t
posterior_variance = _extract_into_tensor(self.posterior_variance, t,
x_t.shape)
posterior_log_variance_clipped = _extract_into_tensor(self.
posterior_log_variance_clipped, t, x_t.shape)
assert posterior_mean.shape[0] == posterior_variance.shape[0
] == posterior_log_variance_clipped.shape[0] == x_start.shape[0]
return posterior_mean, posterior_variance, posterior_log_variance_clipped
",point_e\diffusion\gaussian_diffusion.py
p_mean_variance,"Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
the initial x, x_0.
:param model: the model, which takes a signal and a batch of timesteps
as input.
:param x: the [N x C x ...] tensor at time t.
:param t: a 1-D Tensor of timesteps.
:param clip_denoised: if True, clip the denoised signal into [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample. Applies before
clip_denoised.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict with the following keys:
- 'mean': the model mean output.
- 'variance': the model variance output.
- 'log_variance': the log of 'variance'.
- 'pred_xstart': the prediction for x_0.","def p_mean_variance(self, model, x, t, clip_denoised=False, denoised_fn=
None, model_kwargs=None):
""""""""""""
if model_kwargs is None:
model_kwargs = {}
B, C = x.shape[:2]
assert t.shape == (B,)
model_output = model(x, t, **model_kwargs)
if isinstance(model_output, tuple):
model_output, extra = model_output
else:
extra = None
if self.model_var_type in ['learned', 'learned_range']:
assert model_output.shape == (B, C * 2, *x.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
if self.model_var_type == 'learned':
model_log_variance = model_var_values
model_variance = th.exp(model_log_variance)
else:
min_log = _extract_into_tensor(self.
posterior_log_variance_clipped, t, x.shape)
max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
frac = (model_var_values + 1) / 2
model_log_variance = frac * max_log + (1 - frac) * min_log
model_variance = th.exp(model_log_variance)
else:
model_variance, model_log_variance = {'fixed_large': (np.append(
self.posterior_variance[1], self.betas[1:]), np.log(np.append(
self.posterior_variance[1], self.betas[1:]))), 'fixed_small': (
self.posterior_variance, self.posterior_log_variance_clipped)}[self
.model_var_type]
model_variance = _extract_into_tensor(model_variance, t, x.shape)
model_log_variance = _extract_into_tensor(model_log_variance, t, x.
shape)
def process_xstart(x):
if denoised_fn is not None:
x = denoised_fn(x)
if clip_denoised:
return x.clamp(-1, 1)
return x
if self.model_mean_type == 'x_prev':
pred_xstart = process_xstart(self._predict_xstart_from_xprev(x_t=x,
t=t, xprev=model_output))
model_mean = model_output
elif self.model_mean_type in ['x_start', 'epsilon']:
if self.model_mean_type == 'x_start':
pred_xstart = process_xstart(model_output)
else:
pred_xstart = process_xstart(self._predict_xstart_from_eps(x_t=
x, t=t, eps=model_output))
model_mean, _, _ = self.q_posterior_mean_variance(x_start=
pred_xstart, x_t=x, t=t)
else:
raise NotImplementedError(self.model_mean_type)
assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
return {'mean': model_mean, 'variance': model_variance, 'log_variance':
model_log_variance, 'pred_xstart': pred_xstart, 'extra': extra}
",point_e\diffusion\gaussian_diffusion.py
_predict_xstart_from_eps,,"def _predict_xstart_from_eps(self, x_t, t, eps):
assert x_t.shape == eps.shape
return _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape
) * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,
x_t.shape) * eps
",point_e\diffusion\gaussian_diffusion.py
_predict_xstart_from_xprev,,"def _predict_xstart_from_xprev(self, x_t, t, xprev):
assert x_t.shape == xprev.shape
return _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape
) * xprev - _extract_into_tensor(self.posterior_mean_coef2 / self.
posterior_mean_coef1, t, x_t.shape) * x_t
",point_e\diffusion\gaussian_diffusion.py
_predict_eps_from_xstart,,"def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.
shape) * x_t - pred_xstart) / _extract_into_tensor(self.
sqrt_recipm1_alphas_cumprod, t, x_t.shape)
",point_e\diffusion\gaussian_diffusion.py
condition_mean,"Compute the mean for the previous step, given a function cond_fn that
computes the gradient of a conditional log probability with respect to
x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
condition on y.
This uses the conditioning strategy from Sohl-Dickstein et al. (2015).","def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
""""""""""""
gradient = cond_fn(x, t, **model_kwargs)
new_mean = p_mean_var['mean'].float() + p_mean_var['variance'
] * gradient.float()
return new_mean
",point_e\diffusion\gaussian_diffusion.py
condition_score,"Compute what the p_mean_variance output would have been, should the
model's score function be conditioned by cond_fn.
See condition_mean() for details on cond_fn.
Unlike condition_mean(), this instead uses the conditioning strategy
from Song et al (2020).","def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
""""""""""""
alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])
eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
out = p_mean_var.copy()
out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)
out['mean'], _, _ = self.q_posterior_mean_variance(x_start=out[
'pred_xstart'], x_t=x, t=t)
return out
",point_e\diffusion\gaussian_diffusion.py
p_sample,"Sample x_{t-1} from the model at the given timestep.
:param model: the model to sample from.
:param x: the current tensor at x_{t-1}.
:param t: the value of t, starting at 0 for the first diffusion step.
:param clip_denoised: if True, clip the x_start prediction to [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample.
:param cond_fn: if not None, this is a gradient function that acts
similarly to the model.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict containing the following keys:
- 'sample': a random sample from the model.
- 'pred_xstart': a prediction of x_0.","def p_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
cond_fn=None, model_kwargs=None):
""""""""""""
out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
denoised_fn=denoised_fn, model_kwargs=model_kwargs)
noise = th.randn_like(x)
nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
if cond_fn is not None:
out['mean'] = self.condition_mean(cond_fn, out, x, t, model_kwargs=
model_kwargs)
sample = out['mean'] + nonzero_mask * th.exp(0.5 * out['log_variance']
) * noise
return {'sample': sample, 'pred_xstart': out['pred_xstart']}
",point_e\diffusion\gaussian_diffusion.py
p_sample_loop,"Generate samples from the model.
:param model: the model module.
:param shape: the shape of the samples, (N, C, H, W).
:param noise: if specified, the noise from the encoder to sample.
Should be of the same shape as `shape`.
:param clip_denoised: if True, clip x_start predictions to [-1, 1].
:param denoised_fn: if not None, a function which applies to the
x_start prediction before it is used to sample.
:param cond_fn: if not None, this is a gradient function that acts
similarly to the model.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:param device: if specified, the device to create the samples on.
If not specified, use a model parameter's device.
:param progress: if True, show a tqdm progress bar.
:return: a non-differentiable batch of samples.","def p_sample_loop(self, model, shape, noise=None, clip_denoised=False,
denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
progress=False, temp=1.0):
""""""""""""
final = None
for sample in self.p_sample_loop_progressive(model, shape, noise=noise,
clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=
cond_fn, model_kwargs=model_kwargs, device=device, progress=
progress, temp=temp):
final = sample
return final['sample']
",point_e\diffusion\gaussian_diffusion.py
p_sample_loop_progressive,"Generate samples from the model and yield intermediate samples from
each timestep of diffusion.
Arguments are the same as p_sample_loop().
Returns a generator over dicts, where each dict is the return value of
p_sample().","def p_sample_loop_progressive(self, model, shape, noise=None, clip_denoised
=False, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
progress=False, temp=1.0):
""""""""""""
if device is None:
device = next(model.parameters()).device
assert isinstance(shape, (tuple, list))
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device) * temp
indices = list(range(self.num_timesteps))[::-1]
if progress:
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
out = self.p_sample(model, img, t, clip_denoised=clip_denoised,
denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=
model_kwargs)
yield self.unscale_out_dict(out)
img = out['sample']
",point_e\diffusion\gaussian_diffusion.py
ddim_sample,"Sample x_{t-1} from the model using DDIM.
Same usage as p_sample().","def ddim_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
cond_fn=None, model_kwargs=None, eta=0.0):
""""""""""""
out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
denoised_fn=denoised_fn, model_kwargs=model_kwargs)
if cond_fn is not None:
out = self.condition_score(cond_fn, out, x, t, model_kwargs=
model_kwargs)
eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])
alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(
1 - alpha_bar / alpha_bar_prev)
noise = th.randn_like(x)
mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_prev) + th.sqrt(1 -
alpha_bar_prev - sigma ** 2) * eps
nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
sample = mean_pred + nonzero_mask * sigma * noise
return {'sample': sample, 'pred_xstart': out['pred_xstart']}
",point_e\diffusion\gaussian_diffusion.py
ddim_reverse_sample,Sample x_{t+1} from the model using DDIM reverse ODE.,"def ddim_reverse_sample(self, model, x, t, clip_denoised=False, denoised_fn
=None, cond_fn=None, model_kwargs=None, eta=0.0):
""""""""""""
assert eta == 0.0, 'Reverse ODE only for deterministic path'
out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
denoised_fn=denoised_fn, model_kwargs=model_kwargs)
if cond_fn is not None:
out = self.condition_score(cond_fn, out, x, t, model_kwargs=
model_kwargs)
eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) *
x - out['pred_xstart']) / _extract_into_tensor(self.
sqrt_recipm1_alphas_cumprod, t, x.shape)
alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_next) + th.sqrt(1 -
alpha_bar_next) * eps
return {'sample': mean_pred, 'pred_xstart': out['pred_xstart']}
",point_e\diffusion\gaussian_diffusion.py
ddim_sample_loop,"Generate samples from the model using DDIM.
Same usage as p_sample_loop().","def ddim_sample_loop(self, model, shape, noise=None, clip_denoised=False,
denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
progress=False, eta=0.0, temp=1.0):
""""""""""""
final = None
for sample in self.ddim_sample_loop_progressive(model, shape, noise=
noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn,
cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress
=progress, eta=eta, temp=temp):
final = sample
return final['sample']
",point_e\diffusion\gaussian_diffusion.py
ddim_sample_loop_progressive,"Use DDIM to sample from the model and yield intermediate samples from
each timestep of DDIM.
Same usage as p_sample_loop_progressive().","def ddim_sample_loop_progressive(self, model, shape, noise=None,
clip_denoised=False, denoised_fn=None, cond_fn=None, model_kwargs=None,
device=None, progress=False, eta=0.0, temp=1.0):
""""""""""""
if device is None:
device = next(model.parameters()).device
assert isinstance(shape, (tuple, list))
if noise is not None:
img = noise
else:
img = th.randn(*shape, device=device) * temp
indices = list(range(self.num_timesteps))[::-1]
if progress:
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
t = th.tensor([i] * shape[0], device=device)
with th.no_grad():
out = self.ddim_sample(model, img, t, clip_denoised=
clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn,
model_kwargs=model_kwargs, eta=eta)
yield self.unscale_out_dict(out)
img = out['sample']
",point_e\diffusion\gaussian_diffusion.py
_vb_terms_bpd,"Get a term for the variational lower-bound.
The resulting units are bits (rather than nats, as one might expect).
This allows for comparison to other papers.
:return: a dict with the following keys:
- 'output': a shape [N] tensor of NLLs or KLs.
- 'pred_xstart': the x_0 predictions.","def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False,
model_kwargs=None):
""""""""""""
true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
x_start=x_start, x_t=x_t, t=t)
out = self.p_mean_variance(model, x_t, t, clip_denoised=clip_denoised,
model_kwargs=model_kwargs)
kl = normal_kl(true_mean, true_log_variance_clipped, out['mean'], out[
'log_variance'])
kl = mean_flat(kl) / np.log(2.0)
decoder_nll = -discretized_gaussian_log_likelihood(x_start, means=out[
'mean'], log_scales=0.5 * out['log_variance'])
if not self.discretized_t0:
decoder_nll = th.zeros_like(decoder_nll)
assert decoder_nll.shape == x_start.shape
decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
output = th.where(t == 0, decoder_nll, kl)
return {'output': output, 'pred_xstart': out['pred_xstart'], 'extra':
out['extra']}
",point_e\diffusion\gaussian_diffusion.py
training_losses,"Compute training losses for a single timestep.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param t: a batch of timestep indices.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:param noise: if specified, the specific Gaussian noise to try to remove.
:return: a dict with the key ""loss"" containing a tensor of shape [N].
Some mean or variance settings may also have other keys.","def training_losses(self, model, x_start, t, model_kwargs=None, noise=None
) ->Dict[str, th.Tensor]:
""""""""""""
x_start = self.scale_channels(x_start)
if model_kwargs is None:
model_kwargs = {}
if noise is None:
noise = th.randn_like(x_start)
x_t = self.q_sample(x_start, t, noise=noise)
terms = {}
if self.loss_type == 'kl' or self.loss_type == 'rescaled_kl':
vb_terms = self._vb_terms_bpd(model=model, x_start=x_start, x_t=x_t,
t=t, clip_denoised=False, model_kwargs=model_kwargs)
terms['loss'] = vb_terms['output']
if self.loss_type == 'rescaled_kl':
terms['loss'] *= self.num_timesteps
extra = vb_terms['extra']
elif self.loss_type == 'mse' or self.loss_type == 'rescaled_mse':
model_output = model(x_t, t, **model_kwargs)
if isinstance(model_output, tuple):
model_output, extra = model_output
else:
extra = {}
if self.model_var_type in ['learned', 'learned_range']:
B, C = x_t.shape[:2]
assert model_output.shape == (B, C * 2, *x_t.shape[2:])
model_output, model_var_values = th.split(model_output, C, dim=1)
frozen_out = th.cat([model_output.detach(), model_var_values],
dim=1)
terms['vb'] = self._vb_terms_bpd(model=lambda *args, r=
frozen_out: r, x_start=x_start, x_t=x_t, t=t, clip_denoised
=False)['output']
if self.loss_type == 'rescaled_mse':
terms['vb'] *= self.num_timesteps / 1000.0
target = {'x_prev': self.q_posterior_mean_variance(x_start=x_start,
x_t=x_t, t=t)[0], 'x_start': x_start, 'epsilon': noise}[self.
model_mean_type]
assert model_output.shape == target.shape == x_start.shape
terms['mse'] = mean_flat((target - model_output) ** 2)
if 'vb' in terms:
terms['loss'] = terms['mse'] + terms['vb']
else:
terms['loss'] = terms['mse']
else:
raise NotImplementedError(self.loss_type)
if 'losses' in extra:
terms.update({k: loss for k, (loss, _scale) in extra['losses'].items()}
)
for loss, scale in extra['losses'].values():
terms['loss'] = terms['loss'] + loss * scale
return terms
",point_e\diffusion\gaussian_diffusion.py
_prior_bpd,"Get the prior KL term for the variational lower-bound, measured in
bits-per-dim.
This term can't be optimized, as it only depends on the encoder.
:param x_start: the [N x C x ...] tensor of inputs.
:return: a batch of [N] KL values (in bits), one per batch element.","def _prior_bpd(self, x_start):
""""""""""""
batch_size = x_start.shape[0]
t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0,
logvar2=0.0)
return mean_flat(kl_prior) / np.log(2.0)
",point_e\diffusion\gaussian_diffusion.py
calc_bpd_loop,"Compute the entire variational lower-bound, measured in bits-per-dim,
as well as other related quantities.
:param model: the model to evaluate loss on.
:param x_start: the [N x C x ...] tensor of inputs.
:param clip_denoised: if True, clip denoised samples.
:param model_kwargs: if not None, a dict of extra keyword arguments to
pass to the model. This can be used for conditioning.
:return: a dict containing the following keys:
- total_bpd: the total variational lower-bound, per batch element.
- prior_bpd: the prior term in the lower-bound.
- vb: an [N x T] tensor of terms in the lower-bound.
- xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
- mse: an [N x T] tensor of epsilon MSEs for each timestep.","def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None
):
""""""""""""
device = x_start.device
batch_size = x_start.shape[0]
vb = []
xstart_mse = []
mse = []
for t in list(range(self.num_timesteps))[::-1]:
t_batch = th.tensor([t] * batch_size, device=device)
noise = th.randn_like(x_start)
x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
with th.no_grad():
out = self._vb_terms_bpd(model, x_start=x_start, x_t=x_t, t=
t_batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs
)
vb.append(out['output'])
xstart_mse.append(mean_flat((out['pred_xstart'] - x_start) ** 2))
eps = self._predict_eps_from_xstart(x_t, t_batch, out['pred_xstart'])
mse.append(mean_flat((eps - noise) ** 2))
vb = th.stack(vb, dim=1)
xstart_mse = th.stack(xstart_mse, dim=1)
mse = th.stack(mse, dim=1)
prior_bpd = self._prior_bpd(x_start)
total_bpd = vb.sum(dim=1) + prior_bpd
return {'total_bpd': total_bpd, 'prior_bpd': prior_bpd, 'vb': vb,
'xstart_mse': xstart_mse, 'mse': mse}
",point_e\diffusion\gaussian_diffusion.py
scale_channels,,"def scale_channels(self, x: th.Tensor) ->th.Tensor:
if self.channel_scales is not None:
x = x * th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
[1] * (len(x.shape) - 2))])
if self.channel_biases is not None:
x = x + th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
[1] * (len(x.shape) - 2))])
return x
",point_e\diffusion\gaussian_diffusion.py
unscale_channels,,"def unscale_channels(self, x: th.Tensor) ->th.Tensor:
if self.channel_biases is not None:
x = x - th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
[1] * (len(x.shape) - 2))])
if self.channel_scales is not None:
x = x / th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
[1] * (len(x.shape) - 2))])
return x
",point_e\diffusion\gaussian_diffusion.py
unscale_out_dict,,"def unscale_out_dict(self, out: Dict[str, Union[th.Tensor, Any]]) ->Dict[
str, Union[th.Tensor, Any]]:
return {k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v
) for k, v in out.items()}
",point_e\diffusion\gaussian_diffusion.py
__init__,,"def __init__(self, use_timesteps: Iterable[int], **kwargs):
self.use_timesteps = set(use_timesteps)
self.timestep_map = []
self.original_num_steps = len(kwargs['betas'])
base_diffusion = GaussianDiffusion(**kwargs)
last_alpha_cumprod = 1.0
new_betas = []
for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
if i in self.use_timesteps:
new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
last_alpha_cumprod = alpha_cumprod
self.timestep_map.append(i)
kwargs['betas'] = np.array(new_betas)
super().__init__(**kwargs)
",point_e\diffusion\gaussian_diffusion.py
p_mean_variance,,"def p_mean_variance(self, model, *args, **kwargs):
return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
",point_e\diffusion\gaussian_diffusion.py
training_losses,,"def training_losses(self, model, *args, **kwargs):
return super().training_losses(self._wrap_model(model), *args, **kwargs)
",point_e\diffusion\gaussian_diffusion.py
condition_mean,,"def condition_mean(self, cond_fn, *args, **kwargs):
return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
",point_e\diffusion\gaussian_diffusion.py
condition_score,,"def condition_score(self, cond_fn, *args, **kwargs):
return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
",point_e\diffusion\gaussian_diffusion.py
_wrap_model,,"def _wrap_model(self, model):
if isinstance(model, _WrappedModel):
return model
return _WrappedModel(model, self.timestep_map, self.original_num_steps)
",point_e\diffusion\gaussian_diffusion.py
__init__,,"def __init__(self, model, timestep_map, original_num_steps):
self.model = model
self.timestep_map = timestep_map
self.original_num_steps = original_num_steps
",point_e\diffusion\gaussian_diffusion.py
__call__,,"def __call__(self, x, ts, **kwargs):
map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
new_ts = map_tensor[ts]
return self.model(x, new_ts, **kwargs)
",point_e\diffusion\gaussian_diffusion.py
karras_sample,,"def karras_sample(*args, **kwargs):
last = None
for x in karras_sample_progressive(*args, **kwargs):
last = x['x']
return last
",point_e\diffusion\k_diffusion.py
karras_sample_progressive,,"def karras_sample_progressive(diffusion, model, shape, steps, clip_denoised
=True, progress=False, model_kwargs=None, device=None, sigma_min=0.002,
sigma_max=80, rho=7.0, sampler='heun', s_churn=0.0, s_tmin=0.0, s_tmax=
float('inf'), s_noise=1.0, guidance_scale=0.0):
sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
x_T = th.randn(*shape, device=device) * sigma_max
sample_fn = {'heun': sample_heun, 'dpm': sample_dpm, 'ancestral':
sample_euler_ancestral}[sampler]
if sampler != 'ancestral':
sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax,
s_noise=s_noise)
else:
sampler_args = {}
if isinstance(diffusion, KarrasDenoiser):
def denoiser(x_t, sigma):
_, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
if clip_denoised:
denoised = denoised.clamp(-1, 1)
return denoised
elif isinstance(diffusion, GaussianDiffusion):
model = GaussianToKarrasDenoiser(model, diffusion)
def denoiser(x_t, sigma):
_, denoised = model.denoise(x_t, sigma, clip_denoised=
clip_denoised, model_kwargs=model_kwargs)
return denoised
else:
raise NotImplementedError
if guidance_scale != 0 and guidance_scale != 1:
def guided_denoiser(x_t, sigma):
x_t = th.cat([x_t, x_t], dim=0)
sigma = th.cat([sigma, sigma], dim=0)
x_0 = denoiser(x_t, sigma)
cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)
x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)
return x_0
else:
guided_denoiser = denoiser
for obj in sample_fn(guided_denoiser, x_T, sigmas, progress=progress,
**sampler_args):
if isinstance(diffusion, GaussianDiffusion):
yield diffusion.unscale_out_dict(obj)
else:
yield obj
",point_e\diffusion\k_diffusion.py
get_sigmas_karras,Constructs the noise schedule of Karras et al. (2022).,"def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device='cpu'):
""""""""""""
ramp = th.linspace(0, 1, n)
min_inv_rho = sigma_min ** (1 / rho)
max_inv_rho = sigma_max ** (1 / rho)
sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
return append_zero(sigmas).to(device)
",point_e\diffusion\k_diffusion.py
to_d,Converts a denoiser output to a Karras ODE derivative.,"def to_d(x, sigma, denoised):
""""""""""""
return (x - denoised) / append_dims(sigma, x.ndim)
",point_e\diffusion\k_diffusion.py
get_ancestral_step,"Calculates the noise level (sigma_down) to step down to and the amount
of noise to add (sigma_up) when doing an ancestral sampling step.","def get_ancestral_step(sigma_from, sigma_to):
""""""""""""
sigma_up = (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) /
sigma_from ** 2) ** 0.5
sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
return sigma_down, sigma_up
",point_e\diffusion\k_diffusion.py
sample_euler_ancestral,Ancestral sampling with Euler method steps.,"@th.no_grad()
def sample_euler_ancestral(model, x, sigmas, progress=False):
""""""""""""
s_in = x.new_ones([x.shape[0]])
indices = range(len(sigmas) - 1)
if progress:
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
denoised = model(x, sigmas[i] * s_in)
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i],
'pred_xstart': denoised}
d = to_d(x, sigmas[i], denoised)
dt = sigma_down - sigmas[i]
x = x + d * dt
x = x + th.randn_like(x) * sigma_up
yield {'x': x, 'pred_xstart': x}
",point_e\diffusion\k_diffusion.py
sample_heun,Implements Algorithm 2 (Heun steps) from Karras et al. (2022).,"@th.no_grad()
def sample_heun(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=
0.0, s_tmax=float('inf'), s_noise=1.0):
""""""""""""
s_in = x.new_ones([x.shape[0]])
indices = range(len(sigmas) - 1)
if progress:
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
) if s_tmin <= sigmas[i] <= s_tmax else 0.0
eps = th.randn_like(x) * s_noise
sigma_hat = sigmas[i] * (gamma + 1)
if gamma > 0:
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
denoised = denoiser(x, sigma_hat * s_in)
d = to_d(x, sigma_hat, denoised)
yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
'pred_xstart': denoised}
dt = sigmas[i + 1] - sigma_hat
if sigmas[i + 1] == 0:
x = x + d * dt
else:
x_2 = x + d * dt
denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
d_prime = (d + d_2) / 2
x = x + d_prime * dt
yield {'x': x, 'pred_xstart': denoised}
",point_e\diffusion\k_diffusion.py
sample_dpm,A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).,"@th.no_grad()
def sample_dpm(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=0.0,
s_tmax=float('inf'), s_noise=1.0):
""""""""""""
s_in = x.new_ones([x.shape[0]])
indices = range(len(sigmas) - 1)
if progress:
from tqdm.auto import tqdm
indices = tqdm(indices)
for i in indices:
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
) if s_tmin <= sigmas[i] <= s_tmax else 0.0
eps = th.randn_like(x) * s_noise
sigma_hat = sigmas[i] * (gamma + 1)
if gamma > 0:
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
denoised = denoiser(x, sigma_hat * s_in)
d = to_d(x, sigma_hat, denoised)
yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
'denoised': denoised}
sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2
) ** 3
dt_1 = sigma_mid - sigma_hat
dt_2 = sigmas[i + 1] - sigma_hat
x_2 = x + d * dt_1
denoised_2 = denoiser(x_2, sigma_mid * s_in)
d_2 = to_d(x_2, sigma_mid, denoised_2)
x = x + d_2 * dt_2
yield {'x': x, 'pred_xstart': denoised}
",point_e\diffusion\k_diffusion.py
append_dims,Appends dimensions to the end of a tensor until it has target_dims dimensions.,"def append_dims(x, target_dims):
""""""""""""
dims_to_append = target_dims - x.ndim
if dims_to_append < 0:
raise ValueError(
f'input has {x.ndim} dims but target_dims is {target_dims}, which is less'
)
return x[(...,) + (None,) * dims_to_append]
",point_e\diffusion\k_diffusion.py
append_zero,,"def append_zero(x):
return th.cat([x, x.new_zeros([1])])
",point_e\diffusion\k_diffusion.py
__init__,,"def __init__(self, sigma_data: float=0.5):
self.sigma_data = sigma_data
",point_e\diffusion\k_diffusion.py
get_snr,,"def get_snr(self, sigmas):
return sigmas ** -2
",point_e\diffusion\k_diffusion.py
get_sigmas,,"def get_sigmas(self, sigmas):
return sigmas
",point_e\diffusion\k_diffusion.py
get_scalings,,"def get_scalings(self, sigma):
c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)
c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2
) ** 0.5
c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
return c_skip, c_out, c_in
",point_e\diffusion\k_diffusion.py
training_losses,,"def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None
):
if model_kwargs is None:
model_kwargs = {}
if noise is None:
noise = th.randn_like(x_start)
terms = {}
dims = x_start.ndim
x_t = x_start + noise * append_dims(sigmas, dims)
c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)
]
model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)
target = (x_start - c_skip * x_t) / c_out
terms['mse'] = mean_flat((model_output - target) ** 2)
terms['xs_mse'] = mean_flat((denoised - x_start) ** 2)
if 'vb' in terms:
terms['loss'] = terms['mse'] + terms['vb']
else:
terms['loss'] = terms['mse']
return terms
",point_e\diffusion\k_diffusion.py
denoise,,"def denoise(self, model, x_t, sigmas, **model_kwargs):
c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.
get_scalings(sigmas)]
rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
model_output = model(c_in * x_t, rescaled_t, **model_kwargs)
denoised = c_out * model_output + c_skip * x_t
return model_output, denoised
",point_e\diffusion\k_diffusion.py
__init__,,"def __init__(self, model, diffusion):
from scipy import interpolate
self.model = model
self.diffusion = diffusion
self.alpha_cumprod_to_t = interpolate.interp1d(diffusion.alphas_cumprod,
np.arange(0, diffusion.num_timesteps))
",point_e\diffusion\k_diffusion.py
sigma_to_t,,"def sigma_to_t(self, sigma):
alpha_cumprod = 1.0 / (sigma ** 2 + 1)
if alpha_cumprod > self.diffusion.alphas_cumprod[0]:
return 0
elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:
return self.diffusion.num_timesteps - 1
else:
return float(self.alpha_cumprod_to_t(alpha_cumprod))
",point_e\diffusion\k_diffusion.py
denoise,,"def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None):
t = th.tensor([self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()
], dtype=th.long, device=sigmas.device)
c_in = append_dims(1.0 / (sigmas ** 2 + 1) ** 0.5, x_t.ndim)
out = self.diffusion.p_mean_variance(self.model, x_t * c_in, t,
clip_denoised=clip_denoised, model_kwargs=model_kwargs)
return None, out['pred_xstart']
",point_e\diffusion\k_diffusion.py
__init__,,"def __init__(self, device: torch.device, models: Sequence[nn.Module],
diffusions: Sequence[GaussianDiffusion], num_points: Sequence[int],
aux_channels: Sequence[str], model_kwargs_key_filter: Sequence[str]=(
'*',), guidance_scale: Sequence[float]=(3.0, 3.0), clip_denoised: bool=
True, use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[
int]=(64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max:
Sequence[float]=(120, 160), s_churn: Sequence[float]=(3, 0)):
n = len(models)
assert n > 0
if n > 1:
if len(guidance_scale) == 1:
guidance_scale = list(guidance_scale) + [1.0] * (n - 1)
if len(use_karras) == 1:
use_karras = use_karras * n
if len(karras_steps) == 1:
karras_steps = karras_steps * n
if len(sigma_min) == 1:
sigma_min = sigma_min * n
if len(sigma_max) == 1:
sigma_max = sigma_max * n
if len(s_churn) == 1:
s_churn = s_churn * n
if len(model_kwargs_key_filter) == 1:
model_kwargs_key_filter = model_kwargs_key_filter * n
if len(model_kwargs_key_filter) == 0:
model_kwargs_key_filter = ['*'] * n
assert len(guidance_scale) == n
assert len(use_karras) == n
assert len(karras_steps) == n
assert len(sigma_min) == n
assert len(sigma_max) == n
assert len(s_churn) == n
assert len(model_kwargs_key_filter) == n
self.device = device
self.num_points = num_points
self.aux_channels = aux_channels
self.model_kwargs_key_filter = model_kwargs_key_filter
self.guidance_scale = guidance_scale
self.clip_denoised = clip_denoised
self.use_karras = use_karras
self.karras_steps = karras_steps
self.sigma_min = sigma_min
self.sigma_max = sigma_max
self.s_churn = s_churn
self.models = models
self.diffusions = diffusions
",point_e\diffusion\sampler.py
num_stages,,"@property
def num_stages(self) ->int:
return len(self.models)
",point_e\diffusion\sampler.py
sample_batch,,"def sample_batch(self, batch_size: int, model_kwargs: Dict[str, Any]
) ->torch.Tensor:
samples = None
for x in self.sample_batch_progressive(batch_size, model_kwargs):
samples = x
return samples
",point_e\diffusion\sampler.py
sample_batch_progressive,,"def sample_batch_progressive(self, batch_size: int, model_kwargs: Dict[str,
Any]) ->Iterator[torch.Tensor]:
samples = None
for model, diffusion, stage_num_points, stage_guidance_scale, stage_use_karras, stage_karras_steps, stage_sigma_min, stage_sigma_max, stage_s_churn, stage_key_filter in zip(
self.models, self.diffusions, self.num_points, self.guidance_scale,
self.use_karras, self.karras_steps, self.sigma_min, self.sigma_max,
self.s_churn, self.model_kwargs_key_filter):
stage_model_kwargs = model_kwargs.copy()
if stage_key_filter != '*':
use_keys = set(stage_key_filter.split(','))
stage_model_kwargs = {k: v for k, v in stage_model_kwargs.items
() if k in use_keys}
if samples is not None:
stage_model_kwargs['low_res'] = samples
if hasattr(model, 'cached_model_kwargs'):
stage_model_kwargs = model.cached_model_kwargs(batch_size,
stage_model_kwargs)
sample_shape = batch_size, 3 + len(self.aux_channels), stage_num_points
if stage_guidance_scale != 1 and stage_guidance_scale != 0:
for k, v in stage_model_kwargs.copy().items():
stage_model_kwargs[k] = torch.cat([v, torch.zeros_like(v)],
dim=0)
if stage_use_karras:
samples_it = karras_sample_progressive(diffusion=diffusion,
model=model, shape=sample_shape, steps=stage_karras_steps,
clip_denoised=self.clip_denoised, model_kwargs=
stage_model_kwargs, device=self.device, sigma_min=
stage_sigma_min, sigma_max=stage_sigma_max, s_churn=
stage_s_churn, guidance_scale=stage_guidance_scale)
else:
internal_batch_size = batch_size
if stage_guidance_scale:
model = self._uncond_guide_model(model, stage_guidance_scale)
internal_batch_size *= 2
samples_it = diffusion.p_sample_loop_progressive(model, shape=(
internal_batch_size, *sample_shape[1:]), model_kwargs=
stage_model_kwargs, device=self.device, clip_denoised=self.
clip_denoised)
for x in samples_it:
samples = x['pred_xstart'][:batch_size]
if 'low_res' in stage_model_kwargs:
samples = torch.cat([stage_model_kwargs['low_res'][:len(
samples)], samples], dim=-1)
yield samples
",point_e\diffusion\sampler.py
combine,,"@classmethod
def combine(cls, *samplers: 'PointCloudSampler') ->'PointCloudSampler':
assert all(x.device == samplers[0].device for x in samplers[1:])
assert all(x.aux_channels == samplers[0].aux_channels for x in samplers[1:]
)
assert all(x.clip_denoised == samplers[0].clip_denoised for x in
samplers[1:])
return cls(device=samplers[0].device, models=[x for y in samplers for x in
y.models], diffusions=[x for y in samplers for x in y.diffusions],
num_points=[x for y in samplers for x in y.num_points],
aux_channels=samplers[0].aux_channels, model_kwargs_key_filter=[x for
y in samplers for x in y.model_kwargs_key_filter], guidance_scale=[
x for y in samplers for x in y.guidance_scale], clip_denoised=
samplers[0].clip_denoised, use_karras=[x for y in samplers for x in
y.use_karras], karras_steps=[x for y in samplers for x in y.
karras_steps], sigma_min=[x for y in samplers for x in y.sigma_min],
sigma_max=[x for y in samplers for x in y.sigma_max], s_churn=[x for
y in samplers for x in y.s_churn])
",point_e\diffusion\sampler.py
_uncond_guide_model,,"def _uncond_guide_model(self, model: Callable[..., torch.Tensor], scale: float
) ->Callable[..., torch.Tensor]:
def model_fn(x_t, ts, **kwargs):
half = x_t[:len(x_t) // 2]
combined = torch.cat([half, half], dim=0)
model_out = model(combined, ts, **kwargs)
eps, rest = model_out[:, :3], model_out[:, 3:]
cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
eps = torch.cat([half_eps, half_eps], dim=0)
return torch.cat([eps, rest], dim=1)
return model_fn
",point_e\diffusion\sampler.py
split_model_output,,"def split_model_output(self, output: torch.Tensor, rescale_colors: bool=False
) ->Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
assert len(self.aux_channels) + 3 == output.shape[1
], 'there must be three spatial channels before aux'
pos, joined_aux = output[:, :3], output[:, 3:]
aux = {}
for i, name in enumerate(self.aux_channels):
v = joined_aux[:, i]
if name in {'R', 'G', 'B', 'A'}:
v = v.clamp(0, 255).round()
if rescale_colors:
v = v / 255.0
aux[name] = v
return pos, aux
",point_e\diffusion\sampler.py
output_to_point_clouds,,"def output_to_point_clouds(self, output: torch.Tensor) ->List[PointCloud]:
res = []
for sample in output:
xyz, aux = self.split_model_output(sample[None], rescale_colors=True)
res.append(PointCloud(coords=xyz[0].t().cpu().numpy(), channels={k:
v[0].cpu().numpy() for k, v in aux.items()}))
return res
",point_e\diffusion\sampler.py
with_options,,"def with_options(self, guidance_scale: float, clip_denoised: bool,
use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[int]=(
64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max: Sequence
[float]=(120, 160), s_churn: Sequence[float]=(3, 0)) ->'PointCloudSampler':
return PointCloudSampler(device=self.device, models=self.models,
diffusions=self.diffusions, num_points=self.num_points,
aux_channels=self.aux_channels, model_kwargs_key_filter=self.
model_kwargs_key_filter, guidance_scale=guidance_scale,
clip_denoised=clip_denoised, use_karras=use_karras, karras_steps=
karras_steps, sigma_min=sigma_min, sigma_max=sigma_max, s_churn=s_churn
)
",point_e\diffusion\sampler.py
get_torch_devices,,"def get_torch_devices() ->List[Union[str, torch.device]]:
if torch.cuda.is_available():
return [torch.device(f'cuda:{i}') for i in range(torch.cuda.
device_count())]
else:
return ['cpu']
",point_e\evals\feature_extractor.py
normalize_point_clouds,,"def normalize_point_clouds(pc: np.ndarray) ->np.ndarray:
centroids = np.mean(pc, axis=1, keepdims=True)
pc = pc - centroids
m = np.max(np.sqrt(np.sum(pc ** 2, axis=-1, keepdims=True)), axis=1,
keepdims=True)
pc = pc / m
return pc
",point_e\evals\feature_extractor.py
supports_predictions,,"@property
@abstractmethod
def supports_predictions(self) ->bool:
pass
",point_e\evals\feature_extractor.py
feature_dim,,"@property
@abstractmethod
def feature_dim(self) ->int:
pass
",point_e\evals\feature_extractor.py
num_classes,,"@property
@abstractmethod
def num_classes(self) ->int:
pass
",point_e\evals\feature_extractor.py
features_and_preds,"For a stream of point cloud batches, compute feature vectors and class
predictions.
:param point_clouds: a streamer for a sample batch. Typically, arr_0
will contain the XYZ coordinates.
:return: a tuple (features, predictions)
- features: a [B x feature_dim] array of feature vectors.
- predictions: a [B x num_classes] array of probabilities.","@abstractmethod
def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
ndarray]:
""""""""""""
",point_e\evals\feature_extractor.py
__init__,,"def __init__(self, devices: List[Union[str, torch.device]],
device_batch_size: int=64, cache_dir: Optional[str]=None):
state_dict = load_checkpoint('pointnet', device=torch.device('cpu'),
cache_dir=cache_dir)['model_state_dict']
self.device_batch_size = device_batch_size
self.devices = devices
self.models = []
for device in devices:
model = get_model(num_class=40, normal_channel=False, width_mult=2)
model.load_state_dict(state_dict)
model.to(device)
model.eval()
self.models.append(model)
",point_e\evals\feature_extractor.py
supports_predictions,,"@property
def supports_predictions(self) ->bool:
return True
",point_e\evals\feature_extractor.py
feature_dim,,"@property
def feature_dim(self) ->int:
return 256
",point_e\evals\feature_extractor.py
num_classes,,"@property
def num_classes(self) ->int:
return 40
",point_e\evals\feature_extractor.py
features_and_preds,,"def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
ndarray]:
batch_size = self.device_batch_size * len(self.devices)
point_clouds = (x['arr_0'] for x in streamer.stream(batch_size, ['arr_0']))
output_features = []
output_predictions = []
with ThreadPool(len(self.devices)) as pool:
for batch in point_clouds:
batch = normalize_point_clouds(batch)
batches = []
for i, device in zip(range(0, len(batch), self.
device_batch_size), self.devices):
batches.append(torch.from_numpy(batch[i:i + self.
device_batch_size]).permute(0, 2, 1).to(dtype=torch.
float32, device=device))
def compute_features(i_batch):
i, batch = i_batch
with torch.no_grad():
return self.models[i](batch, features=True)
for logits, _, features in pool.imap(compute_features,
enumerate(batches)):
output_features.append(features.cpu().numpy())
output_predictions.append(logits.exp().cpu().numpy())
return np.concatenate(output_features, axis=0), np.concatenate(
output_predictions, axis=0)
",point_e\evals\feature_extractor.py
compute_statistics,,"def compute_statistics(feats: np.ndarray) ->FIDStatistics:
mu = np.mean(feats, axis=0)
sigma = np.cov(feats, rowvar=False)
return FIDStatistics(mu, sigma)
",point_e\evals\fid_is.py
compute_inception_score,,"def compute_inception_score(preds: np.ndarray, split_size: int=5000) ->float:
scores = []
for i in range(0, len(preds), split_size):
part = preds[i:i + split_size]
kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))
)
kl = np.mean(np.sum(kl, 1))
scores.append(np.exp(kl))
return float(np.mean(scores))
",point_e\evals\fid_is.py
__init__,,"def __init__(self, mu: np.ndarray, sigma: np.ndarray):
self.mu = mu
self.sigma = sigma
",point_e\evals\fid_is.py
frechet_distance,Compute the Frechet distance between two sets of statistics.,"def frechet_distance(self, other, eps=1e-06):
""""""""""""
mu1, sigma1 = self.mu, self.sigma
mu2, sigma2 = other.mu, other.sigma
mu1 = np.atleast_1d(mu1)
mu2 = np.atleast_1d(mu2)
sigma1 = np.atleast_2d(sigma1)
sigma2 = np.atleast_2d(sigma2)
assert mu1.shape == mu2.shape, f'Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}'
assert sigma1.shape == sigma2.shape, f'Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}'
diff = mu1 - mu2
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
if not np.isfinite(covmean).all():
msg = (
'fid calculation produces singular product; adding %s to diagonal of cov estimates'
% eps)
warnings.warn(msg)
offset = np.eye(sigma1.shape[0]) * eps
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
if np.iscomplexobj(covmean):
if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
m = np.max(np.abs(covmean.imag))
raise ValueError('Imaginary component {}'.format(m))
covmean = covmean.real
tr_covmean = np.trace(covmean)
return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2
) - 2 * tr_covmean
",point_e\evals\fid_is.py
_npz_paths_and_length,,"def _npz_paths_and_length(glob_path: str) ->Tuple[List[str], Optional[int]]:
count_match = re.match('^(.*)\\[:([0-9]*)\\]$', glob_path)
if count_match:
raw_path = count_match[1]
max_count = int(count_match[2])
else:
raw_path = glob_path
max_count = None
paths = sorted(glob.glob(raw_path))
if not len(paths):
raise ValueError(f'no paths found matching: {glob_path}')
return paths, max_count
",point_e\evals\npz_stream.py
open_npz_arrays,,"@contextmanager
def open_npz_arrays(path: str, arr_names: Sequence[str]) ->List[NpzArrayReader
]:
if not len(arr_names):
yield []
return
arr_name = arr_names[0]
with open_array(path, arr_name) as arr_f:
version = np.lib.format.read_magic(arr_f)
header = None
if version == (1, 0):
header = np.lib.format.read_array_header_1_0(arr_f)
elif version == (2, 0):
header = np.lib.format.read_array_header_2_0(arr_f)
if header is None:
reader = MemoryNpzArrayReader.load(path, arr_name)
else:
shape, fortran, dtype = header
if fortran or dtype.hasobject:
reader = MemoryNpzArrayReader.load(path, arr_name)
else:
reader = StreamingNpzArrayReader(arr_f, shape, dtype)
with open_npz_arrays(path, arr_names[1:]) as next_readers:
yield [reader] + next_readers
",point_e\evals\npz_stream.py
_read_bytes,"Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886
Read from file-like object until size bytes are read.
Raises ValueError if not EOF is encountered before size bytes are read.
Non-blocking objects only supported if they derive from io objects.
Required as e.g. ZipExtFile in python 2.6 can return less data than
requested.","def _read_bytes(fp, size, error_template='ran out of data'):
""""""""""""
data = bytes()
while True:
try:
r = fp.read(size - len(data))
data += r
if len(r) == 0 or len(data) == size:
break
except io.BlockingIOError:
pass
if len(data) != size:
msg = 'EOF: reading %s, expected %d bytes got %d'
raise ValueError(msg % (error_template, size, len(data)))
else:
return data
",point_e\evals\npz_stream.py
open_array,,"@contextmanager
def open_array(path: str, arr_name: str):
with open(path, 'rb') as f:
with zipfile.ZipFile(f, 'r') as zip_f:
if f'{arr_name}.npy' not in zip_f.namelist():
raise ValueError(f'missing {arr_name} in npz file')
with zip_f.open(f'{arr_name}.npy', 'r') as arr_f:
yield arr_f
",point_e\evals\npz_stream.py
_dict_batch_size,,"def _dict_batch_size(objs: Dict[str, np.ndarray]) ->int:
return len(next(iter(objs.values())))
",point_e\evals\npz_stream.py
infos_from_first_file,,"@classmethod
def infos_from_first_file(cls, glob_path: str) ->Dict[str, 'NumpyArrayInfo']:
paths, _ = _npz_paths_and_length(glob_path)
return cls.infos_from_file(paths[0])
",point_e\evals\npz_stream.py
infos_from_file,Extract the info of every array in an npz file.,"@classmethod
def infos_from_file(cls, npz_path: str) ->Dict[str, 'NumpyArrayInfo']:
""""""""""""
if not os.path.exists(npz_path):
raise FileNotFoundError(f'batch of samples was not found: {npz_path}')
results = {}
with open(npz_path, 'rb') as f:
with zipfile.ZipFile(f, 'r') as zip_f:
for name in zip_f.namelist():
if not name.endswith('.npy'):
continue
key_name = name[:-len('.npy')]
with zip_f.open(name, 'r') as arr_f:
version = np.lib.format.read_magic(arr_f)
if version == (1, 0):
header = np.lib.format.read_array_header_1_0(arr_f)
elif version == (2, 0):
header = np.lib.format.read_array_header_2_0(arr_f)
else:
raise ValueError(
f'unknown numpy array version: {version}')
shape, _, dtype = header
results[key_name] = cls(name=key_name, dtype=dtype,
shape=shape)
return results
",point_e\evals\npz_stream.py
elem_shape,,"@property
def elem_shape(self) ->Tuple[int]:
return self.shape[1:]
",point_e\evals\npz_stream.py
validate,,"def validate(self):
if self.name in {'R', 'G', 'B'}:
if len(self.shape) != 2:
raise ValueError(
f""expecting exactly 2-D shape for '{self.name}' but got: {self.shape}""
)
elif self.name == 'arr_0':
if len(self.shape) < 2:
raise ValueError(
f'expecting at least 2-D shape but got: {self.shape}')
elif len(self.shape) == 3:
if not np.issubdtype(self.dtype, np.floating):
raise ValueError(
f'invalid dtype for audio batch: {self.dtype} (expected float)'
)
elif self.dtype != np.uint8:
raise ValueError(
f'invalid dtype for image batch: {self.dtype} (expected uint8)'
)
",point_e\evals\npz_stream.py
__init__,,"def __init__(self, glob_path: str):
self.paths, self.trunc_length = _npz_paths_and_length(glob_path)
self.infos = NumpyArrayInfo.infos_from_file(self.paths[0])
",point_e\evals\npz_stream.py
keys,,"def keys(self) ->List[str]:
return list(self.infos.keys())
",point_e\evals\npz_stream.py
stream,,"def stream(self, batch_size: int, keys: Sequence[str]) ->Iterator[Dict[str,
np.ndarray]]:
cur_batch = None
num_remaining = self.trunc_length
for path in self.paths:
if num_remaining is not None and num_remaining <= 0:
break
with open_npz_arrays(path, keys) as readers:
combined_reader = CombinedReader(keys, readers)
while num_remaining is None or num_remaining > 0:
read_bs = batch_size
if cur_batch is not None:
read_bs -= _dict_batch_size(cur_batch)
if num_remaining is not None:
read_bs = min(read_bs, num_remaining)
batch = combined_reader.read_batch(read_bs)
if batch is None:
break
if num_remaining is not None:
num_remaining -= _dict_batch_size(batch)
if cur_batch is None:
cur_batch = batch
else:
cur_batch = {k: np.concatenate([cur_batch[k], v], axis=
0) for k, v in batch.items()}
if _dict_batch_size(cur_batch) == batch_size:
yield cur_batch
cur_batch = None
if cur_batch is not None:
yield cur_batch
",point_e\evals\npz_stream.py
read_batch,,"@abstractmethod
def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
pass
",point_e\evals\npz_stream.py
__init__,,"def __init__(self, arr_f, shape, dtype):
self.arr_f = arr_f
self.shape = shape
self.dtype = dtype
self.idx = 0
",point_e\evals\npz_stream.py
read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
if self.idx >= self.shape[0]:
return None
bs = min(batch_size, self.shape[0] - self.idx)
self.idx += bs
if self.dtype.itemsize == 0:
return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
read_count = bs * np.prod(self.shape[1:])
read_size = int(read_count * self.dtype.itemsize)
data = _read_bytes(self.arr_f, read_size, 'array data')
return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
",point_e\evals\npz_stream.py
__init__,,"def __init__(self, arr):
self.arr = arr
self.idx = 0
",point_e\evals\npz_stream.py
load,,"@classmethod
def load(cls, path: str, arr_name: str):
with open(path, 'rb') as f:
arr = np.load(f)[arr_name]
return cls(arr)
",point_e\evals\npz_stream.py
read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
if self.idx >= self.arr.shape[0]:
return None
res = self.arr[self.idx:self.idx + batch_size]
self.idx += batch_size
return res
",point_e\evals\npz_stream.py
__init__,,"def __init__(self, keys: List[str], readers: List[NpzArrayReader]):
self.keys = keys
self.readers = readers
",point_e\evals\npz_stream.py
read_batch,,"def read_batch(self, batch_size: int) ->Optional[Dict[str, np.ndarray]]:
batches = [r.read_batch(batch_size) for r in self.readers]
any_none = any(x is None for x in batches)
all_none = all(x is None for x in batches)
if any_none != all_none:
raise RuntimeError('different keys had different numbers of elements')
if any_none:
return None
if any(len(x) != len(batches[0]) for x in batches):
raise RuntimeError('different keys had different numbers of elements')
return dict(zip(self.keys, batches))
",point_e\evals\npz_stream.py
__init__,,"def __init__(self, num_class, normal_channel=True, width_mult=1):
super(get_model, self).__init__()
self.width_mult = width_mult
in_channel = 6 if normal_channel else 3
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=32,
in_channel=in_channel, mlp=[64 * width_mult, 64 * width_mult, 128 *
width_mult], group_all=False)
self.sa2 = PointNetSetAbstraction(npoint=128, radius=0.4, nsample=64,
in_channel=128 * width_mult + 3, mlp=[128 * width_mult, 128 *
width_mult, 256 * width_mult], group_all=False)
self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=
None, in_channel=256 * width_mult + 3, mlp=[256 * width_mult, 512 *
width_mult, 1024 * width_mult], group_all=True)
self.fc1 = nn.Linear(1024 * width_mult, 512 * width_mult)
self.bn1 = nn.BatchNorm1d(512 * width_mult)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512 * width_mult, 256 * width_mult)
self.bn2 = nn.BatchNorm1d(256 * width_mult)
self.drop2 = nn.Dropout(0.4)
self.fc3 = nn.Linear(256 * width_mult, num_class)
",point_e\evals\pointnet2_cls_ssg.py
forward,,"def forward(self, xyz, features=False):
B, _, _ = xyz.shape
if self.normal_channel:
norm = xyz[:, 3:, :]
xyz = xyz[:, :3, :]
else:
norm = None
l1_xyz, l1_points = self.sa1(xyz, norm)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
x = l3_points.view(B, 1024 * self.width_mult)
x = self.drop1(F.relu(self.bn1(self.fc1(x))))
result_features = self.bn2(self.fc2(x))
x = self.drop2(F.relu(result_features))
x = self.fc3(x)
x = F.log_softmax(x, -1)
if features:
return x, l3_points, result_features
else:
return x, l3_points
",point_e\evals\pointnet2_cls_ssg.py
__init__,,"def __init__(self):
super(get_loss, self).__init__()
",point_e\evals\pointnet2_cls_ssg.py
forward,,"def forward(self, pred, target, trans_feat):
total_loss = F.nll_loss(pred, target)
return total_loss
",point_e\evals\pointnet2_cls_ssg.py
timeit,,"def timeit(tag, t):
print('{}: {}s'.format(tag, time() - t))
return time()
",point_e\evals\pointnet2_utils.py
pc_normalize,,"def pc_normalize(pc):
l = pc.shape[0]
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
pc = pc / m
return pc
",point_e\evals\pointnet2_utils.py
square_distance,"Calculate Euclid distance between each two points.
src^T * dst = xn * xm + yn * ym + zn * zm;
sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
= sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
Input:
src: source points, [B, N, C]
dst: target points, [B, M, C]
Output:
dist: per-point square distance, [B, N, M]","def square_distance(src, dst):
""""""""""""
B, N, _ = src.shape
_, M, _ = dst.shape
dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
dist += torch.sum(src ** 2, -1).view(B, N, 1)
dist += torch.sum(dst ** 2, -1).view(B, 1, M)
return dist
",point_e\evals\pointnet2_utils.py
index_points,"Input:
points: input points data, [B, N, C]
idx: sample index data, [B, S]
Return:
new_points:, indexed points data, [B, S, C]","def index_points(points, idx):
""""""""""""
device = points.device
B = points.shape[0]
view_shape = list(idx.shape)
view_shape[1:] = [1] * (len(view_shape) - 1)
repeat_shape = list(idx.shape)
repeat_shape[0] = 1
batch_indices = torch.arange(B, dtype=torch.long).to(device).view(
view_shape).repeat(repeat_shape)
new_points = points[batch_indices, idx, :]
return new_points
",point_e\evals\pointnet2_utils.py
farthest_point_sample,"Input:
xyz: pointcloud data, [B, N, 3]
npoint: number of samples
Return:
centroids: sampled pointcloud index, [B, npoint]","def farthest_point_sample(xyz, npoint, deterministic=False):
""""""""""""
device = xyz.device
B, N, C = xyz.shape
centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
distance = torch.ones(B, N).to(device) * 10000000000.0
if deterministic:
farthest = torch.arange(0, B, dtype=torch.long).to(device)
else:
farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
batch_indices = torch.arange(B, dtype=torch.long).to(device)
for i in range(npoint):
centroids[:, i] = farthest
centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
dist = torch.sum((xyz - centroid) ** 2, -1)
mask = dist < distance
distance[mask] = dist[mask]
farthest = torch.max(distance, -1)[1]
return centroids
",point_e\evals\pointnet2_utils.py
query_ball_point,"Input:
radius: local region radius
nsample: max sample number in local region
xyz: all points, [B, N, 3]
new_xyz: query points, [B, S, 3]
Return:
group_idx: grouped points index, [B, S, nsample]","def query_ball_point(radius, nsample, xyz, new_xyz):
""""""""""""
device = xyz.device
B, N, C = xyz.shape
_, S, _ = new_xyz.shape
group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N
).repeat([B, S, 1])
sqrdists = square_distance(new_xyz, xyz)
group_idx[sqrdists > radius ** 2] = N
group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
mask = group_idx == N
group_idx[mask] = group_first[mask]
return group_idx
",point_e\evals\pointnet2_utils.py
sample_and_group,"Input:
npoint:
radius:
nsample:
xyz: input points position data, [B, N, 3]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, npoint, nsample, 3]
new_points: sampled points data, [B, npoint, nsample, 3+D]","def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False,
deterministic=False):
""""""""""""
B, N, C = xyz.shape
S = npoint
fps_idx = farthest_point_sample(xyz, npoint, deterministic=deterministic)
new_xyz = index_points(xyz, fps_idx)
idx = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = index_points(xyz, idx)
grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
if points is not None:
grouped_points = index_points(points, idx)
new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1)
else:
new_points = grouped_xyz_norm
if returnfps:
return new_xyz, new_points, grouped_xyz, fps_idx
else:
return new_xyz, new_points
",point_e\evals\pointnet2_utils.py
sample_and_group_all,"Input:
xyz: input points position data, [B, N, 3]
points: input points data, [B, N, D]
Return:
new_xyz: sampled points position data, [B, 1, 3]
new_points: sampled points data, [B, 1, N, 3+D]","def sample_and_group_all(xyz, points):
""""""""""""
device = xyz.device
B, N, C = xyz.shape
new_xyz = torch.zeros(B, 1, C).to(device)
grouped_xyz = xyz.view(B, 1, N, C)
if points is not None:
new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
else:
new_points = grouped_xyz
return new_xyz, new_points
",point_e\evals\pointnet2_utils.py
__init__,,"def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
super(PointNetSetAbstraction, self).__init__()
self.npoint = npoint
self.radius = radius
self.nsample = nsample
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.group_all = group_all
",point_e\evals\pointnet2_utils.py
forward,"Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
""""""""""""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
if self.group_all:
new_xyz, new_points = sample_and_group_all(xyz, points)
else:
new_xyz, new_points = sample_and_group(self.npoint, self.radius,
self.nsample, xyz, points, deterministic=not self.training)
new_points = new_points.permute(0, 3, 2, 1)
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
new_points = torch.max(new_points, 2)[0]
new_xyz = new_xyz.permute(0, 2, 1)
return new_xyz, new_points
",point_e\evals\pointnet2_utils.py
__init__,,"def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
super(PointNetSetAbstractionMsg, self).__init__()
self.npoint = npoint
self.radius_list = radius_list
self.nsample_list = nsample_list
self.conv_blocks = nn.ModuleList()
self.bn_blocks = nn.ModuleList()
for i in range(len(mlp_list)):
convs = nn.ModuleList()
bns = nn.ModuleList()
last_channel = in_channel + 3
for out_channel in mlp_list[i]:
convs.append(nn.Conv2d(last_channel, out_channel, 1))
bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.conv_blocks.append(convs)
self.bn_blocks.append(bns)
",point_e\evals\pointnet2_utils.py
forward,"Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
""""""""""""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
B, N, C = xyz.shape
S = self.npoint
new_xyz = index_points(xyz, farthest_point_sample(xyz, S, deterministic
=not self.training))
new_points_list = []
for i, radius in enumerate(self.radius_list):
K = self.nsample_list[i]
group_idx = query_ball_point(radius, K, xyz, new_xyz)
grouped_xyz = index_points(xyz, group_idx)
grouped_xyz -= new_xyz.view(B, S, 1, C)
if points is not None:
grouped_points = index_points(points, group_idx)
grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
else:
grouped_points = grouped_xyz
grouped_points = grouped_points.permute(0, 3, 2, 1)
for j in range(len(self.conv_blocks[i])):
conv = self.conv_blocks[i][j]
bn = self.bn_blocks[i][j]
grouped_points = F.relu(bn(conv(grouped_points)))
new_points = torch.max(grouped_points, 2)[0]
new_points_list.append(new_points)
new_xyz = new_xyz.permute(0, 2, 1)
new_points_concat = torch.cat(new_points_list, dim=1)
return new_xyz, new_points_concat
",point_e\evals\pointnet2_utils.py
__init__,,"def __init__(self, in_channel, mlp):
super(PointNetFeaturePropagation, self).__init__()
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm1d(out_channel))
last_channel = out_channel
",point_e\evals\pointnet2_utils.py
forward,"Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]","def forward(self, xyz1, xyz2, points1, points2):
""""""""""""
xyz1 = xyz1.permute(0, 2, 1)
xyz2 = xyz2.permute(0, 2, 1)
points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
if S == 1:
interpolated_points = points2.repeat(1, N, 1)
else:
dists = square_distance(xyz1, xyz2)
dists, idx = dists.sort(dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3]
dist_recip = 1.0 / (dists + 1e-08)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
interpolated_points = torch.sum(index_points(points2, idx) * weight
.view(B, N, 3, 1), dim=2)
if points1 is not None:
points1 = points1.permute(0, 2, 1)
new_points = torch.cat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
return new_points
",point_e\evals\pointnet2_utils.py
clear_scene,,"def clear_scene():
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()
",point_e\evals\scripts\blender_script.py
clear_lights,,"def clear_lights():
bpy.ops.object.select_all(action='DESELECT')
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, bpy.types.Light):
obj.select_set(True)
bpy.ops.object.delete()
",point_e\evals\scripts\blender_script.py
import_model,,"def import_model(path):
clear_scene()
_, ext = os.path.splitext(path)
ext = ext.lower()
if ext == '.obj':
bpy.ops.import_scene.obj(filepath=path)
elif ext in ['.glb', '.gltf']:
bpy.ops.import_scene.gltf(filepath=path)
elif ext == '.stl':
bpy.ops.import_mesh.stl(filepath=path)
elif ext == '.fbx':
bpy.ops.import_scene.fbx(filepath=path)
elif ext == '.dae':
bpy.ops.wm.collada_import(filepath=path)
elif ext == '.ply':
bpy.ops.import_mesh.ply(filepath=path)
else:
raise RuntimeError(f'unexpected extension: {ext}')
",point_e\evals\scripts\blender_script.py
scene_root_objects,,"def scene_root_objects():
for obj in bpy.context.scene.objects.values():
if not obj.parent:
yield obj
",point_e\evals\scripts\blender_script.py
scene_bbox,,"def scene_bbox(single_obj=None, ignore_matrix=False):
bbox_min = (math.inf,) * 3
bbox_max = (-math.inf,) * 3
found = False
for obj in (scene_meshes() if single_obj is None else [single_obj]):
found = True
for coord in obj.bound_box:
coord = Vector(coord)
if not ignore_matrix:
coord = obj.matrix_world @ coord
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
if not found:
raise RuntimeError('no objects in scene to compute bounding box for')
return Vector(bbox_min), Vector(bbox_max)
",point_e\evals\scripts\blender_script.py
scene_meshes,,"def scene_meshes():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, bpy.types.Mesh):
yield obj
",point_e\evals\scripts\blender_script.py
normalize_scene,,"def normalize_scene():
bbox_min, bbox_max = scene_bbox()
scale = 1 / max(bbox_max - bbox_min)
for obj in scene_root_objects():
obj.scale = obj.scale * scale
bpy.context.view_layer.update()
bbox_min, bbox_max = scene_bbox()
offset = -(bbox_min + bbox_max) / 2
for obj in scene_root_objects():
obj.matrix_world.translation += offset
bpy.ops.object.select_all(action='DESELECT')
",point_e\evals\scripts\blender_script.py
create_camera,,"def create_camera():
camera_data = bpy.data.cameras.new(name='Camera')
camera_object = bpy.data.objects.new('Camera', camera_data)
bpy.context.scene.collection.objects.link(camera_object)
bpy.context.scene.camera = camera_object
",point_e\evals\scripts\blender_script.py
set_camera,,"def set_camera(direction, camera_dist=2.0):
camera_pos = -camera_dist * direction
bpy.context.scene.camera.location = camera_pos
rot_quat = direction.to_track_quat('-Z', 'Y')
bpy.context.scene.camera.rotation_euler = rot_quat.to_euler()
bpy.context.view_layer.update()
",point_e\evals\scripts\blender_script.py
randomize_camera,,"def randomize_camera(camera_dist=2.0):
direction = random_unit_vector()
set_camera(direction, camera_dist=camera_dist)
",point_e\evals\scripts\blender_script.py
pan_camera,,"def pan_camera(time, axis='Z', camera_dist=2.0, elevation=-0.1):
angle = time * math.pi * 2
direction = [-math.cos(angle), -math.sin(angle), -elevation]
assert axis in ['X', 'Y', 'Z']
if axis == 'X':
direction = [direction[2], *direction[:2]]
elif axis == 'Y':
direction = [direction[0], -elevation, direction[1]]
direction = Vector(direction).normalized()
set_camera(direction, camera_dist=camera_dist)
",point_e\evals\scripts\blender_script.py
place_camera,,"def place_camera(time, camera_pose_mode='random', camera_dist_min=2.0,
camera_dist_max=2.0):
camera_dist = random.uniform(camera_dist_min, camera_dist_max)
if camera_pose_mode == 'random':
randomize_camera(camera_dist=camera_dist)
elif camera_pose_mode == 'z-circular':
pan_camera(time, axis='Z', camera_dist=camera_dist)
elif camera_pose_mode == 'z-circular-elevated':
pan_camera(time, axis='Z', camera_dist=camera_dist, elevation=
0.2617993878)
else:
raise ValueError(f'Unknown camera pose mode: {camera_pose_mode}')
",point_e\evals\scripts\blender_script.py
create_light,,"def create_light(location, energy=1.0, angle=0.5 * math.pi / 180):
light_data = bpy.data.lights.new(name='Light', type='SUN')
light_data.energy = energy
light_data.angle = angle
light_object = bpy.data.objects.new(name='Light', object_data=light_data)
direction = -location
rot_quat = direction.to_track_quat('-Z', 'Y')
light_object.rotation_euler = rot_quat.to_euler()
bpy.context.view_layer.update()
bpy.context.collection.objects.link(light_object)
light_object.location = location
",point_e\evals\scripts\blender_script.py
create_random_lights,,"def create_random_lights(count=4, distance=2.0, energy=1.5):
clear_lights()
for _ in range(count):
create_light(random_unit_vector() * distance, energy=energy)
",point_e\evals\scripts\blender_script.py
create_camera_light,,"def create_camera_light():
clear_lights()
create_light(bpy.context.scene.camera.location, energy=5.0)
",point_e\evals\scripts\blender_script.py
create_uniform_light,,"def create_uniform_light(backend):
clear_lights()
pos = Vector(UNIFORM_LIGHT_DIRECTION)
angle = 0.0092 if backend == 'CYCLES' else math.pi
create_light(pos, energy=5.0, angle=angle)
create_light(-pos, energy=5.0, angle=angle)
",point_e\evals\scripts\blender_script.py
create_vertex_color_shaders,,"def create_vertex_color_shaders():
for obj in bpy.context.scene.objects.values():
if not isinstance(obj.data, bpy.types.Mesh):
continue
if len(obj.data.materials):
continue
color_keys = (obj.data.vertex_colors or {}).keys()
if not len(color_keys):
continue
mat = bpy.data.materials.new(name='VertexColored')
mat.use_nodes = True
bsdf_node = None
for node in mat.node_tree.nodes:
if node.type == 'BSDF_PRINCIPLED':
bsdf_node = node
assert bsdf_node is not None, 'material has no Principled BSDF node to modify'
socket_map = {}
for input in bsdf_node.inputs:
socket_map[input.name] = input
socket_map['Specular'].default_value = 0.0
socket_map['Roughness'].default_value = 1.0
v_color = mat.node_tree.nodes.new('ShaderNodeVertexColor')
v_color.layer_name = color_keys[0]
mat.node_tree.links.new(v_color.outputs[0], socket_map['Base Color'])
obj.data.materials.append(mat)
",point_e\evals\scripts\blender_script.py
create_default_materials,,"def create_default_materials():
for obj in bpy.context.scene.objects.values():
if isinstance(obj.data, bpy.types.Mesh):
if not len(obj.data.materials):
mat = bpy.data.materials.new(name='DefaultMaterial')
mat.use_nodes = True
obj.data.materials.append(mat)
",point_e\evals\scripts\blender_script.py
find_materials,,"def find_materials():
all_materials = set()
for obj in bpy.context.scene.objects.values():
if not isinstance(obj.data, bpy.types.Mesh):
continue
for mat in obj.data.materials:
all_materials.add(mat)
return all_materials
",point_e\evals\scripts\blender_script.py
get_socket_value,,"def get_socket_value(tree, socket):
default = socket.default_value
if not isinstance(default, float):
default = list(default)
for link in tree.links:
if link.to_socket == socket:
return link.from_socket, default
return None, default
",point_e\evals\scripts\blender_script.py
clear_socket_input,,"def clear_socket_input(tree, socket):
for link in list(tree.links):
if link.to_socket == socket:
tree.links.remove(link)
",point_e\evals\scripts\blender_script.py
set_socket_value,,"def set_socket_value(tree, socket, socket_and_default):
clear_socket_input(tree, socket)
old_source_socket, default = socket_and_default
if isinstance(default, float) and not isinstance(socket.default_value,
float):
socket.default_value = [default] * 3 + [1.0]
else:
socket.default_value = default
if old_source_socket is not None:
tree.links.new(old_source_socket, socket)
",point_e\evals\scripts\blender_script.py
setup_nodes,,"def setup_nodes(output_path, capturing_material_alpha: bool=False):
tree = bpy.context.scene.node_tree
links = tree.links
for node in tree.nodes:
tree.nodes.remove(node)
def node_op(op: str, *args, clamp=False):
node = tree.nodes.new(type='CompositorNodeMath')
node.operation = op
if clamp:
node.use_clamp = True
for i, arg in enumerate(args):
if isinstance(arg, (int, float)):
node.inputs[i].default_value = arg
else:
links.new(arg, node.inputs[i])
return node.outputs[0]
def node_clamp(x, maximum=1.0):
return node_op('MINIMUM', x, maximum)
def node_mul(x, y, **kwargs):
return node_op('MULTIPLY', x, y, **kwargs)
input_node = tree.nodes.new(type='CompositorNodeRLayers')
input_node.scene = bpy.context.scene
input_sockets = {}
for output in input_node.outputs:
input_sockets[output.name] = output
if capturing_material_alpha:
color_socket = input_sockets['Image']
else:
raw_color_socket = input_sockets['Image']
color_node = tree.nodes.new(type='CompositorNodeConvertColorSpace')
color_node.from_color_space = 'Linear'
color_node.to_color_space = 'sRGB'
tree.links.new(raw_color_socket, color_node.inputs[0])
color_socket = color_node.outputs[0]
split_node = tree.nodes.new(type='CompositorNodeSepRGBA')
tree.links.new(color_socket, split_node.inputs[0])
for i, channel in (enumerate('rgba') if not capturing_material_alpha else
[(0, 'MatAlpha')]):
output_node = tree.nodes.new(type='CompositorNodeOutputFile')
output_node.base_path = f'{output_path}_{channel}'
links.new(split_node.outputs[i], output_node.inputs[0])
if capturing_material_alpha:
return
depth_out = node_clamp(node_mul(input_sockets['Depth'], 1 / MAX_DEPTH))
output_node = tree.nodes.new(type='CompositorNodeOutputFile')
output_node.base_path = f'{output_path}_depth'
links.new(depth_out, output_node.inputs[0])
",point_e\evals\scripts\blender_script.py
render_scene,,"def render_scene(output_path, fast_mode: bool):
use_workbench = bpy.context.scene.render.engine == 'BLENDER_WORKBENCH'
if use_workbench:
bpy.context.scene.render.engine = 'BLENDER_EEVEE'
bpy.context.scene.eevee.taa_render_samples = 1
if fast_mode:
if bpy.context.scene.render.engine == 'BLENDER_EEVEE':
bpy.context.scene.eevee.taa_render_samples = 1
elif bpy.context.scene.render.engine == 'CYCLES':
bpy.context.scene.cycles.samples = 256
elif bpy.context.scene.render.engine == 'CYCLES':
bpy.context.scene.cycles.time_limit = 40
bpy.context.view_layer.update()
bpy.context.scene.use_nodes = True
bpy.context.scene.view_layers['ViewLayer'].use_pass_z = True
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.film_transparent = True
bpy.context.scene.render.resolution_x = 512
bpy.context.scene.render.resolution_y = 512
bpy.context.scene.render.image_settings.file_format = 'PNG'
bpy.context.scene.render.image_settings.color_mode = 'BW'
bpy.context.scene.render.image_settings.color_depth = '16'
bpy.context.scene.render.filepath = output_path
setup_nodes(output_path)
bpy.ops.render.render(write_still=True)
for channel_name in ['r', 'g', 'b', 'a', 'depth']:
sub_dir = f'{output_path}_{channel_name}'
image_path = os.path.join(sub_dir, os.listdir(sub_dir)[0])
name, ext = os.path.splitext(output_path)
if channel_name == 'depth' or not use_workbench:
os.rename(image_path, f'{name}_{channel_name}{ext}')
else:
os.remove(image_path)
os.removedirs(sub_dir)
if use_workbench:
bpy.context.scene.use_nodes = False
bpy.context.scene.render.engine = 'BLENDER_WORKBENCH'
bpy.context.scene.render.image_settings.color_mode = 'RGBA'
bpy.context.scene.render.image_settings.color_depth = '8'
bpy.context.scene.display.shading.color_type = 'TEXTURE'
bpy.context.scene.display.shading.light = 'FLAT'
if fast_mode:
bpy.context.scene.display.render_aa = 'FXAA'
os.remove(output_path)
bpy.ops.render.render(write_still=True)
bpy.context.scene.render.image_settings.color_mode = 'BW'
bpy.context.scene.render.image_settings.color_depth = '16'
",point_e\evals\scripts\blender_script.py
scene_fov,,"def scene_fov():
x_fov = bpy.context.scene.camera.data.angle_x
y_fov = bpy.context.scene.camera.data.angle_y
width = bpy.context.scene.render.resolution_x
height = bpy.context.scene.render.resolution_y
if bpy.context.scene.camera.data.angle == x_fov:
y_fov = 2 * math.atan(math.tan(x_fov / 2) * height / width)
else:
x_fov = 2 * math.atan(math.tan(y_fov / 2) * width / height)
return x_fov, y_fov
",point_e\evals\scripts\blender_script.py
write_camera_metadata,,"def write_camera_metadata(path):
x_fov, y_fov = scene_fov()
bbox_min, bbox_max = scene_bbox()
matrix = bpy.context.scene.camera.matrix_world
with open(path, 'w') as f:
json.dump(dict(format_version=FORMAT_VERSION, max_depth=MAX_DEPTH,
bbox=[list(bbox_min), list(bbox_max)], origin=list(matrix.col[3
])[:3], x_fov=x_fov, y_fov=y_fov, x=list(matrix.col[0])[:3], y=
list(-matrix.col[1])[:3], z=list(-matrix.col[2])[:3]), f)
",point_e\evals\scripts\blender_script.py
save_rendering_dataset,,"def save_rendering_dataset(input_path: str, output_path: str, num_images:
int, backend: str, light_mode: str, camera_pose: str, camera_dist_min:
float, camera_dist_max: float, fast_mode: bool):
assert light_mode in ['random', 'uniform', 'camera']
assert camera_pose in ['random', 'z-circular', 'z-circular-elevated']
import_model(input_path)
bpy.context.scene.render.engine = backend
normalize_scene()
if light_mode == 'random':
create_random_lights()
elif light_mode == 'uniform':
create_uniform_light(backend)
create_camera()
create_vertex_color_shaders()
for i in range(num_images):
t = i / max(num_images - 1, 1)
place_camera(t, camera_pose_mode=camera_pose, camera_dist_min=
camera_dist_min, camera_dist_max=camera_dist_max)
if light_mode == 'camera':
create_camera_light()
render_scene(os.path.join(output_path, f'{i:05}.png'), fast_mode=
fast_mode)
write_camera_metadata(os.path.join(output_path, f'{i:05}.json'))
with open(os.path.join(output_path, 'info.json'), 'w') as f:
info = dict(backend=backend, light_mode=light_mode, fast_mode=
fast_mode, format_version=FORMAT_VERSION, channels=['R', 'G',
'B', 'A', 'D'], scale=0.5)
json.dump(info, f)
",point_e\evals\scripts\blender_script.py
main,,"def main():
try:
dash_index = sys.argv.index('--')
except ValueError as exc:
raise ValueError(""arguments must be preceded by '--'"") from exc
raw_args = sys.argv[dash_index + 1:]
parser = argparse.ArgumentParser()
parser.add_argument('--input_path', required=True, type=str)
parser.add_argument('--output_path', required=True, type=str)
parser.add_argument('--num_images', type=int, default=20)
parser.add_argument('--backend', type=str, default='BLENDER_EEVEE')
parser.add_argument('--light_mode', type=str, default='uniform')
parser.add_argument('--camera_pose', type=str, default='random')
parser.add_argument('--camera_dist_min', type=float, default=2.0)
parser.add_argument('--camera_dist_max', type=float, default=2.0)
parser.add_argument('--fast_mode', action='store_true')
args = parser.parse_args(raw_args)
save_rendering_dataset(input_path=args.input_path, output_path=args.
output_path, num_images=args.num_images, backend=args.backend,
light_mode=args.light_mode, camera_pose=args.camera_pose,
camera_dist_min=args.camera_dist_min, camera_dist_max=args.
camera_dist_max, fast_mode=args.fast_mode)
",point_e\evals\scripts\blender_script.py
main,,"def main():
parser = argparse.ArgumentParser()
parser.add_argument('--cache_dir', type=str, default=None)
parser.add_argument('batch_1', type=str)
parser.add_argument('batch_2', type=str)
args = parser.parse_args()
print('creating classifier...')
clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
cache_dir)
print('computing first batch activations')
features_1, _ = clf.features_and_preds(NpzStreamer(args.batch_1))
stats_1 = compute_statistics(features_1)
del features_1
features_2, _ = clf.features_and_preds(NpzStreamer(args.batch_2))
stats_2 = compute_statistics(features_2)
del features_2
print(f'P-FID: {stats_1.frechet_distance(stats_2)}')
",point_e\evals\scripts\evaluate_pfid.py
main,,"def main():
parser = argparse.ArgumentParser()
parser.add_argument('--cache_dir', type=str, default=None)
parser.add_argument('batch', type=str)
args = parser.parse_args()
print('creating classifier...')
clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
cache_dir)
print('computing batch predictions')
_, preds = clf.features_and_preds(NpzStreamer(args.batch))
print(f'P-IS: {compute_inception_score(preds)}')
",point_e\evals\scripts\evaluate_pis.py
checkpoint,"Evaluate a function without caching intermediate activations, allowing for
reduced memory at the expense of extra compute in the backward pass.
:param func: the function to evaluate.
:param inputs: the argument sequence to pass to `func`.
:param params: a sequence of parameters `func` depends on but does not
explicitly take as arguments.
:param flag: if False, disable gradient checkpointing.","def checkpoint(func: Callable[..., Union[torch.Tensor, Sequence[torch.
Tensor]]], inputs: Sequence[torch.Tensor], params: Iterable[torch.
Tensor], flag: bool):
""""""""""""
if flag:
args = tuple(inputs) + tuple(params)
return CheckpointFunction.apply(func, len(inputs), *args)
else:
return func(*inputs)
",point_e\models\checkpoint.py
forward,,"@staticmethod
def forward(ctx, run_function, length, *args):
ctx.run_function = run_function
ctx.input_tensors = list(args[:length])
ctx.input_params = list(args[length:])
with torch.no_grad():
output_tensors = ctx.run_function(*ctx.input_tensors)
return output_tensors
",point_e\models\checkpoint.py
backward,,"@staticmethod
def backward(ctx, *output_grads):
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.
input_tensors]
with torch.enable_grad():
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
output_tensors = ctx.run_function(*shallow_copies)
input_grads = torch.autograd.grad(output_tensors, ctx.input_tensors +
ctx.input_params, output_grads, allow_unused=True)
del ctx.input_tensors
del ctx.input_params
del output_tensors
return (None, None) + input_grads
",point_e\models\checkpoint.py
model_from_config,,"def model_from_config(config: Dict[str, Any], device: torch.device
) ->nn.Module:
config = config.copy()
name = config.pop('name')
if name == 'PointDiffusionTransformer':
return PointDiffusionTransformer(device=device, dtype=torch.float32,
**config)
elif name == 'CLIPImagePointDiffusionTransformer':
return CLIPImagePointDiffusionTransformer(device=device, dtype=
torch.float32, **config)
elif name == 'CLIPImageGridPointDiffusionTransformer':
return CLIPImageGridPointDiffusionTransformer(device=device, dtype=
torch.float32, **config)
elif name == 'UpsamplePointDiffusionTransformer':
return UpsamplePointDiffusionTransformer(device=device, dtype=torch
.float32, **config)
elif name == 'CLIPImageGridUpsamplePointDiffusionTransformer':
return CLIPImageGridUpsamplePointDiffusionTransformer(device=device,
dtype=torch.float32, **config)
elif name == 'CrossAttentionPointCloudSDFModel':
return CrossAttentionPointCloudSDFModel(device=device, dtype=torch.
float32, **config)
raise ValueError(f'unknown model name: {name}')
",point_e\models\configs.py
default_cache_dir,,"@lru_cache()
def default_cache_dir() ->str:
return os.path.join(os.path.abspath(os.getcwd()), 'point_e_model_cache')
",point_e\models\download.py
fetch_file_cached,"Download the file at the given URL into a local file and return the path.
If cache_dir is specified, it will be used to download the files.
Otherwise, default_cache_dir() is used.","def fetch_file_cached(url: str, progress: bool=True, cache_dir: Optional[
str]=None, chunk_size: int=4096) ->str:
""""""""""""
if cache_dir is None:
cache_dir = default_cache_dir()
os.makedirs(cache_dir, exist_ok=True)
local_path = os.path.join(cache_dir, url.split('/')[-1])
if os.path.exists(local_path):
return local_path
response = requests.get(url, stream=True)
size = int(response.headers.get('content-length', '0'))
with FileLock(local_path + '.lock'):
if progress:
pbar = tqdm(total=size, unit='iB', unit_scale=True)
tmp_path = local_path + '.tmp'
with open(tmp_path, 'wb') as f:
for chunk in response.iter_content(chunk_size):
if progress:
pbar.update(len(chunk))
f.write(chunk)
os.rename(tmp_path, local_path)
if progress:
pbar.close()
return local_path
",point_e\models\download.py
load_checkpoint,,"def load_checkpoint(checkpoint_name: str, device: torch.device, progress:
bool=True, cache_dir: Optional[str]=None, chunk_size: int=4096) ->Dict[
str, torch.Tensor]:
if checkpoint_name not in MODEL_PATHS:
raise ValueError(
f'Unknown checkpoint name {checkpoint_name}. Known names are: {MODEL_PATHS.keys()}.'
)
path = fetch_file_cached(MODEL_PATHS[checkpoint_name], progress=
progress, cache_dir=cache_dir, chunk_size=chunk_size)
return torch.load(path, map_location=device)
",point_e\models\download.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
width: int, heads: int, init_scale: float, data_width: Optional[int]=None):
super().__init__()
self.n_data = n_data
self.width = width
self.heads = heads
self.data_width = width if data_width is None else data_width
self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
self.c_kv = nn.Linear(self.data_width, width * 2, device=device, dtype=
dtype)
self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
self.attention = QKVMultiheadCrossAttention(device=device, dtype=dtype,
heads=heads, n_data=n_data)
init_linear(self.c_q, init_scale)
init_linear(self.c_kv, init_scale)
init_linear(self.c_proj, init_scale)
",point_e\models\perceiver.py
forward,,"def forward(self, x, data):
x = self.c_q(x)
data = self.c_kv(data)
x = checkpoint(self.attention, (x, data), (), True)
x = self.c_proj(x)
return x
",point_e\models\perceiver.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
n_data: int):
super().__init__()
self.device = device
self.dtype = dtype
self.heads = heads
self.n_data = n_data
",point_e\models\perceiver.py
forward,,"def forward(self, q, kv):
_, n_ctx, _ = q.shape
bs, n_data, width = kv.shape
attn_ch = width // self.heads // 2
scale = 1 / math.sqrt(math.sqrt(attn_ch))
q = q.view(bs, n_ctx, self.heads, -1)
kv = kv.view(bs, n_data, self.heads, -1)
k, v = torch.split(kv, attn_ch, dim=-1)
weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
wdtype = weight.dtype
weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
",point_e\models\perceiver.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
width: int, heads: int, data_width: Optional[int]=None, init_scale:
float=1.0):
super().__init__()
if data_width is None:
data_width = width
self.attn = MultiheadCrossAttention(device=device, dtype=dtype, n_data=
n_data, width=width, heads=heads, data_width=data_width, init_scale
=init_scale)
self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
init_scale)
self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
",point_e\models\perceiver.py
forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
x = x + self.attn(self.ln_1(x), self.ln_2(data))
x = x + self.mlp(self.ln_3(x))
return x
",point_e\models\perceiver.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
width: int, layers: int, heads: int, init_scale: float=0.25, data_width:
Optional[int]=None):
super().__init__()
self.width = width
self.layers = layers
init_scale = init_scale * math.sqrt(1.0 / width)
self.resblocks = nn.ModuleList([ResidualCrossAttentionBlock(device=
device, dtype=dtype, n_data=n_data, width=width, heads=heads,
init_scale=init_scale, data_width=data_width) for _ in range(layers)])
",point_e\models\perceiver.py
forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
for block in self.resblocks:
x = block(x, data)
return x
",point_e\models\perceiver.py
_image_to_pil,,"def _image_to_pil(obj: Optional[ImageType]) ->Image.Image:
if obj is None:
return Image.fromarray(np.zeros([64, 64, 3], dtype=np.uint8))
if isinstance(obj, np.ndarray):
return Image.fromarray(obj.astype(np.uint8))
elif isinstance(obj, torch.Tensor):
return Image.fromarray(obj.detach().cpu().numpy().astype(np.uint8))
else:
return obj
",point_e\models\pretrained_clip.py
__init__,,"def __init__(self, device: torch.device, dtype: Optional[torch.dtype]=torch
.float32, ensure_used_params: bool=True, clip_name: str='ViT-L/14',
cache_dir: Optional[str]=None):
super().__init__()
assert clip_name in ['ViT-L/14', 'ViT-B/32']
self.device = device
self.ensure_used_params = ensure_used_params
import clip
self.clip_model, self.preprocess = clip.load(clip_name, device=device,
download_root=cache_dir or default_cache_dir())
self.clip_name = clip_name
if dtype is not None:
self.clip_model.to(dtype)
self._tokenize = clip.tokenize
",point_e\models\pretrained_clip.py
feature_dim,,"@property
def feature_dim(self) ->int:
if self.clip_name == 'ViT-L/14':
return 768
else:
return 512
",point_e\models\pretrained_clip.py
grid_size,,"@property
def grid_size(self) ->int:
if self.clip_name == 'ViT-L/14':
return 16
else:
return 7
",point_e\models\pretrained_clip.py
grid_feature_dim,,"@property
def grid_feature_dim(self) ->int:
if self.clip_name == 'ViT-L/14':
return 1024
else:
return 768
",point_e\models\pretrained_clip.py
forward,"Generate a batch of embeddings from a mixture of images, texts,
precomputed embeddings, and possibly empty values.
For each batch element, at most one of images, texts, and embeddings
should have a non-None value. Embeddings from multiple modalities
cannot be mixed for a single batch element. If no modality is provided,
a zero embedding will be used for the batch element.","def forward(self, batch_size: int, images: Optional[Iterable[Optional[
ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
) ->torch.Tensor:
""""""""""""
image_seq = [None] * batch_size if images is None else list(images)
text_seq = [None] * batch_size if texts is None else list(texts)
embedding_seq = [None] * batch_size if embeddings is None else list(
embeddings)
assert len(image_seq
) == batch_size, 'number of images should match batch size'
assert len(text_seq
) == batch_size, 'number of texts should match batch size'
assert len(embedding_seq
) == batch_size, 'number of embeddings should match batch size'
if self.ensure_used_params:
return self._static_multimodal_embed(images=image_seq, texts=
text_seq, embeddings=embedding_seq)
result = torch.zeros((batch_size, self.feature_dim), device=self.device)
index_images = []
index_texts = []
for i, (image, text, emb) in enumerate(zip(image_seq, text_seq,
embedding_seq)):
assert sum([int(image is not None), int(text is not None), int(emb
is not None)]
) < 2, 'only one modality may be non-None per batch element'
if image is not None:
index_images.append((i, image))
elif text is not None:
index_texts.append((i, text))
elif emb is not None:
result[i] = emb.to(result)
if len(index_images):
embs = self.embed_images(img for _, img in index_images)
for (i, _), emb in zip(index_images, embs):
result[i] = emb.to(result)
if len(index_texts):
embs = self.embed_text(text for _, text in index_texts)
for (i, _), emb in zip(index_texts, embs):
result[i] = emb.to(result)
return result
",point_e\models\pretrained_clip.py
_static_multimodal_embed,"Like forward(), but always runs all encoders to ensure that
the forward graph looks the same on every rank.","def _static_multimodal_embed(self, images: List[Optional[ImageType]]=None,
texts: List[Optional[str]]=None, embeddings: List[Optional[torch.Tensor
]]=None) ->torch.Tensor:
""""""""""""
image_emb = self.embed_images(images)
text_emb = self.embed_text(t if t else '' for t in texts)
joined_embs = torch.stack([(emb.to(device=self.device, dtype=torch.
float32) if emb is not None else torch.zeros(self.feature_dim,
device=self.device)) for emb in embeddings], dim=0)
image_flag = torch.tensor([(x is not None) for x in images], device=
self.device)[:, None].expand_as(image_emb)
text_flag = torch.tensor([(x is not None) for x in texts], device=self.
device)[:, None].expand_as(image_emb)
emb_flag = torch.tensor([(x is not None) for x in embeddings], device=
self.device)[:, None].expand_as(image_emb)
return image_flag.float() * image_emb + text_flag.float(
) * text_emb + emb_flag.float(
) * joined_embs + self.clip_model.logit_scale * 0
",point_e\models\pretrained_clip.py
embed_images,":param xs: N images, stored as numpy arrays, tensors, or PIL images.
:return: an [N x D] tensor of features.","def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
""""""""""""
clip_inputs = self.images_to_tensor(xs)
results = self.clip_model.encode_image(clip_inputs).float()
return results / torch.linalg.norm(results, dim=-1, keepdim=True)
",point_e\models\pretrained_clip.py
embed_text,Embed text prompts as an [N x D] tensor.,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
""""""""""""
enc = self.clip_model.encode_text(self._tokenize(list(prompts),
truncate=True).to(self.device)).float()
return enc / torch.linalg.norm(enc, dim=-1, keepdim=True)
",point_e\models\pretrained_clip.py
embed_images_grid,"Embed images into latent grids.
:param xs: an iterable of images to embed.
:return: a tensor of shape [N x C x L], where L = self.grid_size**2.","def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
""""""""""""
if self.ensure_used_params:
extra_value = 0.0
for p in self.parameters():
extra_value = extra_value + p.mean() * 0.0
else:
extra_value = 0.0
x = self.images_to_tensor(xs).to(self.clip_model.dtype)
vt = self.clip_model.visual
x = vt.conv1(x)
x = x.reshape(x.shape[0], x.shape[1], -1)
x = x.permute(0, 2, 1)
x = torch.cat([vt.class_embedding.to(x.dtype) + torch.zeros(x.shape[0],
1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)
x = x + vt.positional_embedding.to(x.dtype)
x = vt.ln_pre(x)
x = x.permute(1, 0, 2)
x = vt.transformer(x)
x = x.permute(1, 2, 0)
return x[..., 1:].contiguous().float() + extra_value
",point_e\models\pretrained_clip.py
images_to_tensor,,"def images_to_tensor(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
return torch.stack([self.preprocess(_image_to_pil(x)) for x in xs], dim=0
).to(self.device)
",point_e\models\pretrained_clip.py
__init__,,"def __init__(self, device: torch.device, **kwargs):
self.model = ImageCLIP(device, dtype=None, ensure_used_params=False, **
kwargs)
for parameter in self.model.parameters():
parameter.requires_grad_(False)
",point_e\models\pretrained_clip.py
feature_dim,,"@property
def feature_dim(self) ->int:
return self.model.feature_dim
",point_e\models\pretrained_clip.py
grid_size,,"@property
def grid_size(self) ->int:
return self.model.grid_size
",point_e\models\pretrained_clip.py
grid_feature_dim,,"@property
def grid_feature_dim(self) ->int:
return self.model.grid_feature_dim
",point_e\models\pretrained_clip.py
__call__,,"def __call__(self, batch_size: int, images: Optional[Iterable[Optional[
ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
) ->torch.Tensor:
return self.model(batch_size=batch_size, images=images, texts=texts,
embeddings=embeddings)
",point_e\models\pretrained_clip.py
embed_images,,"def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
with torch.no_grad():
return self.model.embed_images(xs)
",point_e\models\pretrained_clip.py
embed_text,,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
with torch.no_grad():
return self.model.embed_text(prompts)
",point_e\models\pretrained_clip.py
embed_images_grid,,"def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
with torch.no_grad():
return self.model.embed_images_grid(xs)
",point_e\models\pretrained_clip.py
device,Get the device that should be used for input tensors.,"@property
@abstractmethod
def device(self) ->torch.device:
""""""""""""
",point_e\models\sdf.py
default_batch_size,"Get a reasonable default number of query points for the model.
In some cases, this might be the only supported size.","@property
@abstractmethod
def default_batch_size(self) ->int:
""""""""""""
",point_e\models\sdf.py
encode_point_clouds,"Encode a batch of point clouds to cache part of the SDF calculation
done by forward().
:param point_clouds: a batch of [batch x 3 x N] points.
:return: a state representing the encoded point cloud batch.","@abstractmethod
def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
.Tensor]:
""""""""""""
",point_e\models\sdf.py
forward,"Predict the SDF at the coordinates x, given a batch of point clouds.
Either point_clouds or encoded should be passed. Only exactly one of
these arguments should be None.
:param x: a [batch x 3 x N'] tensor of query points.
:param point_clouds: a [batch x 3 x N] batch of point clouds.
:param encoded: the result of calling encode_point_clouds().
:return: a [batch x N'] tensor of SDF predictions.","def forward(self, x: torch.Tensor, point_clouds: Optional[torch.Tensor]=
None, encoded: Optional[Dict[str, torch.Tensor]]=None) ->torch.Tensor:
""""""""""""
assert point_clouds is not None or encoded is not None
assert point_clouds is None or encoded is None
if point_clouds is not None:
encoded = self.encode_point_clouds(point_clouds)
return self.predict_sdf(x, encoded)
",point_e\models\sdf.py
predict_sdf,"Predict the SDF at the query points given the encoded point clouds.
Each query point should be treated independently, only conditioning on
the point clouds themselves.","@abstractmethod
def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
Tensor]]) ->torch.Tensor:
""""""""""""
",point_e\models\sdf.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
4096, width: int=512, encoder_layers: int=12, encoder_heads: int=8,
decoder_layers: int=4, decoder_heads: int=8, init_scale: float=0.25):
super().__init__()
self._device = device
self.n_ctx = n_ctx
self.encoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
self.encoder = Transformer(device=device, dtype=dtype, n_ctx=n_ctx,
width=width, layers=encoder_layers, heads=encoder_heads, init_scale
=init_scale)
self.decoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
self.decoder = SimplePerceiver(device=device, dtype=dtype, n_data=n_ctx,
width=width, layers=decoder_layers, heads=decoder_heads, init_scale
=init_scale)
self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
self.output_proj = nn.Linear(width, 1, device=device, dtype=dtype)
",point_e\models\sdf.py
device,,"@property
def device(self) ->torch.device:
return self._device
",point_e\models\sdf.py
default_batch_size,,"@property
def default_batch_size(self) ->int:
return self.n_query
",point_e\models\sdf.py
encode_point_clouds,,"def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
.Tensor]:
h = self.encoder_input_proj(point_clouds.permute(0, 2, 1))
h = self.encoder(h)
return dict(latents=h)
",point_e\models\sdf.py
predict_sdf,,"def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
Tensor]]) ->torch.Tensor:
data = encoded['latents']
x = self.decoder_input_proj(x.permute(0, 2, 1))
x = self.decoder(x, data)
x = self.ln_post(x)
x = self.output_proj(x)
return x[..., 0]
",point_e\models\sdf.py
init_linear,,"def init_linear(l, stddev):
nn.init.normal_(l.weight, std=stddev)
if l.bias is not None:
nn.init.constant_(l.bias, 0.0)
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
width: int, heads: int, init_scale: float):
super().__init__()
self.n_ctx = n_ctx
self.width = width
self.heads = heads
self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
self.attention = QKVMultiheadAttention(device=device, dtype=dtype,
heads=heads, n_ctx=n_ctx)
init_linear(self.c_qkv, init_scale)
init_linear(self.c_proj, init_scale)
",point_e\models\transformer.py
forward,,"def forward(self, x):
x = self.c_qkv(x)
x = checkpoint(self.attention, (x,), (), True)
x = self.c_proj(x)
return x
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int,
init_scale: float):
super().__init__()
self.width = width
self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
self.gelu = nn.GELU()
init_linear(self.c_fc, init_scale)
init_linear(self.c_proj, init_scale)
",point_e\models\transformer.py
forward,,"def forward(self, x):
return self.c_proj(self.gelu(self.c_fc(x)))
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
n_ctx: int):
super().__init__()
self.device = device
self.dtype = dtype
self.heads = heads
self.n_ctx = n_ctx
",point_e\models\transformer.py
forward,,"def forward(self, qkv):
bs, n_ctx, width = qkv.shape
attn_ch = width // self.heads // 3
scale = 1 / math.sqrt(math.sqrt(attn_ch))
qkv = qkv.view(bs, n_ctx, self.heads, -1)
q, k, v = torch.split(qkv, attn_ch, dim=-1)
weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
wdtype = weight.dtype
weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
width: int, heads: int, init_scale: float=1.0):
super().__init__()
self.attn = MultiheadAttention(device=device, dtype=dtype, n_ctx=n_ctx,
width=width, heads=heads, init_scale=init_scale)
self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
init_scale)
self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)
",point_e\models\transformer.py
forward,,"def forward(self, x: torch.Tensor):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
width: int, layers: int, heads: int, init_scale: float=0.25):
super().__init__()
self.n_ctx = n_ctx
self.width = width
self.layers = layers
init_scale = init_scale * math.sqrt(1.0 / width)
self.resblocks = nn.ModuleList([ResidualAttentionBlock(device=device,
dtype=dtype, n_ctx=n_ctx, width=width, heads=heads, init_scale=
init_scale) for _ in range(layers)])
",point_e\models\transformer.py
forward,,"def forward(self, x: torch.Tensor):
for block in self.resblocks:
x = block(x)
return x
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
input_channels: int=3, output_channels: int=3, n_ctx: int=1024, width:
int=512, layers: int=12, heads: int=8, init_scale: float=0.25,
time_token_cond: bool=False):
super().__init__()
self.input_channels = input_channels
self.output_channels = output_channels
self.n_ctx = n_ctx
self.time_token_cond = time_token_cond
self.time_embed = MLP(device=device, dtype=dtype, width=width,
init_scale=init_scale * math.sqrt(1.0 / width))
self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
self.backbone = Transformer(device=device, dtype=dtype, n_ctx=n_ctx +
int(time_token_cond), width=width, layers=layers, heads=heads,
init_scale=init_scale)
self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
self.input_proj = nn.Linear(input_channels, width, device=device, dtype
=dtype)
self.output_proj = nn.Linear(width, output_channels, device=device,
dtype=dtype)
with torch.no_grad():
self.output_proj.weight.zero_()
self.output_proj.bias.zero_()
",point_e\models\transformer.py
forward,":param x: an [N x C x T] tensor.
:param t: an [N] tensor.
:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor):
""""""""""""
assert x.shape[-1] == self.n_ctx
t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
return self._forward_with_cond(x, [(t_embed, self.time_token_cond)])
",point_e\models\transformer.py
_forward_with_cond,,"def _forward_with_cond(self, x: torch.Tensor, cond_as_token: List[Tuple[
torch.Tensor, bool]]) ->torch.Tensor:
h = self.input_proj(x.permute(0, 2, 1))
for emb, as_token in cond_as_token:
if not as_token:
h = h + emb[:, None]
extra_tokens = [(emb[:, None] if len(emb.shape) == 2 else emb) for emb,
as_token in cond_as_token if as_token]
if len(extra_tokens):
h = torch.cat(extra_tokens + [h], dim=1)
h = self.ln_pre(h)
h = self.backbone(h)
h = self.ln_post(h)
if len(extra_tokens):
h = h[:, sum(h.shape[1] for h in extra_tokens):]
h = self.output_proj(h)
return h.permute(0, 2, 1)
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
1024, token_cond: bool=False, cond_drop_prob: float=0.0, frozen_clip:
bool=True, cache_dir: Optional[str]=None, **kwargs):
super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + int(
token_cond), **kwargs)
self.n_ctx = n_ctx
self.token_cond = token_cond
self.clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
cache_dir=cache_dir)
self.clip_embed = nn.Linear(self.clip.feature_dim, self.backbone.width,
device=device, dtype=dtype)
self.cond_drop_prob = cond_drop_prob
",point_e\models\transformer.py
cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
) ->Dict[str, Any]:
with torch.no_grad():
return dict(embeddings=self.clip(batch_size, **model_kwargs))
",point_e\models\transformer.py
forward,":param x: an [N x C x T] tensor.
:param t: an [N] tensor.
:param images: a batch of images to condition on.
:param texts: a batch of texts to condition on.
:param embeddings: a batch of CLIP embeddings to condition on.
:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
Iterable[Optional[ImageType]]]=None, texts: Optional[Iterable[Optional[
str]]]=None, embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None):
""""""""""""
assert x.shape[-1] == self.n_ctx
t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
clip_out = self.clip(batch_size=len(x), images=images, texts=texts,
embeddings=embeddings)
assert len(clip_out.shape) == 2 and clip_out.shape[0] == x.shape[0]
if self.training:
mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
clip_out = clip_out * mask[:, None].to(clip_out)
clip_out = math.sqrt(clip_out.shape[1]) * clip_out
clip_embed = self.clip_embed(clip_out)
cond = [(clip_embed, self.token_cond), (t_embed, self.time_token_cond)]
return self._forward_with_cond(x, cond)
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
1024, cond_drop_prob: float=0.0, frozen_clip: bool=True, cache_dir:
Optional[str]=None, **kwargs):
clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
cache_dir=cache_dir)
super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
grid_size ** 2, **kwargs)
self.n_ctx = n_ctx
self.clip = clip
self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
self.clip.grid_feature_dim, self.backbone.width, device=device,
dtype=dtype))
self.cond_drop_prob = cond_drop_prob
",point_e\models\transformer.py
cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
) ->Dict[str, Any]:
_ = batch_size
with torch.no_grad():
return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
'images']))
",point_e\models\transformer.py
forward,":param x: an [N x C x T] tensor.
:param t: an [N] tensor.
:param images: a batch of images to condition on.
:param embeddings: a batch of CLIP latent grids to condition on.
:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
Iterable[ImageType]]=None, embeddings: Optional[Iterable[torch.Tensor]]
=None):
""""""""""""
assert images is not None or embeddings is not None, 'must specify images or embeddings'
assert images is None or embeddings is None, 'cannot specify both images and embeddings'
assert x.shape[-1] == self.n_ctx
t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
if images is not None:
clip_out = self.clip.embed_images_grid(images)
else:
clip_out = embeddings
if self.training:
mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
clip_out = clip_out * mask[:, None, None].to(clip_out)
clip_out = clip_out.permute(0, 2, 1)
clip_embed = self.clip_embed(clip_out)
cond = [(t_embed, self.time_token_cond), (clip_embed, True)]
return self._forward_with_cond(x, cond)
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
cond_input_channels: Optional[int]=None, cond_ctx: int=1024, n_ctx: int
=4096 - 1024, channel_scales: Optional[Sequence[float]]=None,
channel_biases: Optional[Sequence[float]]=None, **kwargs):
super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + cond_ctx, **
kwargs)
self.n_ctx = n_ctx
self.cond_input_channels = cond_input_channels or self.input_channels
self.cond_point_proj = nn.Linear(self.cond_input_channels, self.
backbone.width, device=device, dtype=dtype)
self.register_buffer('channel_scales', torch.tensor(channel_scales,
dtype=dtype, device=device) if channel_scales is not None else None)
self.register_buffer('channel_biases', torch.tensor(channel_biases,
dtype=dtype, device=device) if channel_biases is not None else None)
",point_e\models\transformer.py
forward,":param x: an [N x C1 x T] tensor.
:param t: an [N] tensor.
:param low_res: an [N x C2 x T'] tensor of conditioning points.
:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.Tensor):
""""""""""""
assert x.shape[-1] == self.n_ctx
t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
low_res_embed = self._embed_low_res(low_res)
cond = [(t_embed, self.time_token_cond), (low_res_embed, True)]
return self._forward_with_cond(x, cond)
",point_e\models\transformer.py
_embed_low_res,,"def _embed_low_res(self, x: torch.Tensor) ->torch.Tensor:
if self.channel_scales is not None:
x = x * self.channel_scales[None, :, None]
if self.channel_biases is not None:
x = x + self.channel_biases[None, :, None]
return self.cond_point_proj(x.permute(0, 2, 1))
",point_e\models\transformer.py
__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
4096 - 1024, cond_drop_prob: float=0.0, frozen_clip: bool=True,
cache_dir: Optional[str]=None, **kwargs):
clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
cache_dir=cache_dir)
super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
grid_size ** 2, **kwargs)
self.n_ctx = n_ctx
self.clip = clip
self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
self.clip.grid_feature_dim, self.backbone.width, device=device,
dtype=dtype))
self.cond_drop_prob = cond_drop_prob
",point_e\models\transformer.py
cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
) ->Dict[str, Any]:
if 'images' not in model_kwargs:
zero_emb = torch.zeros([batch_size, self.clip.grid_feature_dim,
self.clip.grid_size ** 2], device=next(self.parameters()).device)
return dict(embeddings=zero_emb, low_res=model_kwargs['low_res'])
with torch.no_grad():
return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
'images']), low_res=model_kwargs['low_res'])
",point_e\models\transformer.py
forward,":param x: an [N x C1 x T] tensor.
:param t: an [N] tensor.
:param low_res: an [N x C2 x T'] tensor of conditioning points.
:param images: a batch of images to condition on.
:param embeddings: a batch of CLIP latent grids to condition on.
:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.
Tensor, images: Optional[Iterable[ImageType]]=None, embeddings:
Optional[Iterable[torch.Tensor]]=None):
""""""""""""
assert x.shape[-1] == self.n_ctx
t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
low_res_embed = self._embed_low_res(low_res)
if images is not None:
clip_out = self.clip.embed_images_grid(images)
elif embeddings is not None:
clip_out = embeddings
else:
clip_out = torch.zeros([len(x), self.clip.grid_feature_dim, self.
clip.grid_size ** 2], dtype=x.dtype, device=x.device)
if self.training:
mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
clip_out = clip_out * mask[:, None, None].to(clip_out)
clip_out = clip_out.permute(0, 2, 1)
clip_embed = self.clip_embed(clip_out)
cond = [(t_embed, self.time_token_cond), (clip_embed, True), (
low_res_embed, True)]
return self._forward_with_cond(x, cond)
",point_e\models\transformer.py
timestep_embedding,"Create sinusoidal timestep embeddings.
:param timesteps: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an [N x dim] Tensor of positional embeddings.","def timestep_embedding(timesteps, dim, max_period=10000):
""""""""""""
half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=
half, dtype=torch.float32) / half).to(device=timesteps.device)
args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1]
)], dim=-1)
return embedding
",point_e\models\util.py
load,Load the mesh from a .npz file.,"@classmethod
def load(cls, f: Union[str, BinaryIO]) ->'TriMesh':
""""""""""""
if isinstance(f, str):
with open(f, 'rb') as reader:
return cls.load(reader)
else:
obj = np.load(f)
keys = list(obj.keys())
verts = obj['verts']
faces = obj['faces']
normals = obj['normals'] if 'normals' in keys else None
vertex_channels = {}
face_channels = {}
for key in keys:
if key.startswith('v_'):
vertex_channels[key[2:]] = obj[key]
elif key.startswith('f_'):
face_channels[key[2:]] = obj[key]
return cls(verts=verts, faces=faces, normals=normals,
vertex_channels=vertex_channels, face_channels=face_channels)
",point_e\util\mesh.py
save,Save the mesh to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
""""""""""""
if isinstance(f, str):
with open(f, 'wb') as writer:
self.save(writer)
else:
obj_dict = dict(verts=self.verts, faces=self.faces)
if self.normals is not None:
obj_dict['normals'] = self.normals
for k, v in self.vertex_channels.items():
obj_dict[f'v_{k}'] = v
for k, v in self.face_channels.items():
obj_dict[f'f_{k}'] = v
np.savez(f, **obj_dict)
",point_e\util\mesh.py
has_vertex_colors,,"def has_vertex_colors(self) ->bool:
return self.vertex_channels is not None and all(x in self.
vertex_channels for x in 'RGB')
",point_e\util\mesh.py
write_ply,,"def write_ply(self, raw_f: BinaryIO):
write_ply(raw_f, coords=self.verts, rgb=np.stack([self.vertex_channels[
x] for x in 'RGB'], axis=1) if self.has_vertex_colors() else None,
faces=self.faces)
",point_e\util\mesh.py
marching_cubes_mesh,"Run marching cubes on the SDF predicted from a point cloud to produce a
mesh representing the 3D surface.
:param pc: the point cloud to apply marching cubes to.
:param model: the model to use to predict SDF values.
:param grid_size: the number of samples along each axis. A total of
grid_size**3 function evaluations are performed.
:param side_length: the size of the cube containing the model, which is
assumed to be centered at the origin.
:param fill_vertex_channels: if True, use the nearest neighbor of each mesh
vertex in the point cloud to compute vertex
data (e.g. colors).","def marching_cubes_mesh(pc: PointCloud, model: PointCloudSDFModel,
batch_size: int=4096, grid_size: int=128, side_length: float=1.02,
fill_vertex_channels: bool=True, progress: bool=False) ->TriMesh:
""""""""""""
voxel_size = side_length / (grid_size - 1)
min_coord = -side_length / 2
def int_coord_to_float(int_coords: torch.Tensor) ->torch.Tensor:
return int_coords.float() * voxel_size + min_coord
with torch.no_grad():
cond = model.encode_point_clouds(torch.from_numpy(pc.coords).
permute(1, 0).to(model.device)[None])
indices = range(0, grid_size ** 3, batch_size)
if progress:
indices = tqdm(indices)
volume = []
for i in indices:
indices = torch.arange(i, min(i + batch_size, grid_size ** 3), step
=1, dtype=torch.int64, device=model.device)
zs = int_coord_to_float(indices % grid_size)
ys = int_coord_to_float(torch.div(indices, grid_size, rounding_mode
='trunc') % grid_size)
xs = int_coord_to_float(torch.div(indices, grid_size ** 2,
rounding_mode='trunc'))
coords = torch.stack([xs, ys, zs], dim=0)
with torch.no_grad():
volume.append(model(coords[None], encoded=cond)[0])
volume_np = torch.cat(volume).view(grid_size, grid_size, grid_size).cpu(
).numpy()
if np.all(volume_np < 0) or np.all(volume_np > 0):
volume_np -= np.mean(volume_np)
verts, faces, normals, _ = skimage.measure.marching_cubes(volume=
volume_np, level=0, allow_degenerate=False, spacing=(voxel_size,) * 3)
old_f1 = faces[:, 0].copy()
faces[:, 0] = faces[:, 1]
faces[:, 1] = old_f1
verts += min_coord
return TriMesh(verts=verts, faces=faces, normals=normals,
vertex_channels=None if not fill_vertex_channels else
_nearest_vertex_channels(pc, verts))
",point_e\util\pc_to_mesh.py
_nearest_vertex_channels,,"def _nearest_vertex_channels(pc: PointCloud, verts: np.ndarray) ->Dict[str,
np.ndarray]:
nearest = pc.nearest_points(verts)
return {ch: arr[nearest] for ch, arr in pc.channels.items()}
",point_e\util\pc_to_mesh.py
plot_point_cloud,"Render a point cloud as a plot to the given image path.
:param pc: the PointCloud to plot.
:param image_path: the path to save the image, with a file extension.
:param color: if True, show the RGB colors from the point cloud.
:param grid_size: the number of random rotations to render.","def plot_point_cloud(pc: PointCloud, color: bool=True, grid_size: int=1,
fixed_bounds: Optional[Tuple[Tuple[float, float, float], Tuple[float,
float, float]]]=((-0.75, -0.75, -0.75), (0.75, 0.75, 0.75))):
""""""""""""
fig = plt.figure(figsize=(8, 8))
for i in range(grid_size):
for j in range(grid_size):
ax = fig.add_subplot(grid_size, grid_size, 1 + j + i *
grid_size, projection='3d')
color_args = {}
if color:
color_args['c'] = np.stack([pc.channels['R'], pc.channels[
'G'], pc.channels['B']], axis=-1)
c = pc.coords
if grid_size > 1:
theta = np.pi * 2 * (i * grid_size + j) / grid_size ** 2
rotation = np.array([[np.cos(theta), -np.sin(theta), 0.0],
[np.sin(theta), np.cos(theta), 0.0], [0.0, 0.0, 1.0]])
c = c @ rotation
ax.scatter(c[:, 0], c[:, 1], c[:, 2], **color_args)
if fixed_bounds is None:
min_point = c.min(0)
max_point = c.max(0)
size = (max_point - min_point).max() / 2
center = (min_point + max_point) / 2
ax.set_xlim3d(center[0] - size, center[0] + size)
ax.set_ylim3d(center[1] - size, center[1] + size)
ax.set_zlim3d(center[2] - size, center[2] + size)
else:
ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0])
ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1])
ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2])
return fig
",point_e\util\plotting.py
write_ply,"Write a PLY file for a mesh or a point cloud.
:param coords: an [N x 3] array of floating point coordinates.
:param rgb: an [N x 3] array of vertex colors, in the range [0.0, 1.0].
:param faces: an [N x 3] array of triangles encoded as integer indices.","def write_ply(raw_f: BinaryIO, coords: np.ndarray, rgb: Optional[np.ndarray
]=None, faces: Optional[np.ndarray]=None):
""""""""""""
with buffered_writer(raw_f) as f:
f.write(b'ply\n')
f.write(b'format binary_little_endian 1.0\n')
f.write(bytes(f'element vertex {len(coords)}\n', 'ascii'))
f.write(b'property float x\n')
f.write(b'property float y\n')
f.write(b'property float z\n')
if rgb is not None:
f.write(b'property uchar red\n')
f.write(b'property uchar green\n')
f.write(b'property uchar blue\n')
if faces is not None:
f.write(bytes(f'element face {len(faces)}\n', 'ascii'))
f.write(b'property list uchar int vertex_index\n')
f.write(b'end_header\n')
if rgb is not None:
rgb = (rgb * 255.499).round().astype(int)
vertices = [(*coord, *rgb) for coord, rgb in zip(coords.tolist(
), rgb.tolist())]
format = struct.Struct('<3f3B')
for item in vertices:
f.write(format.pack(*item))
else:
format = struct.Struct('<3f')
for vertex in coords.tolist():
f.write(format.pack(*vertex))
if faces is not None:
format = struct.Struct('<B3I')
for tri in faces.tolist():
f.write(format.pack(len(tri), *tri))
",point_e\util\ply_util.py
buffered_writer,,"@contextmanager
def buffered_writer(raw_f: BinaryIO) ->Iterator[io.BufferedIOBase]:
if isinstance(raw_f, io.BufferedIOBase):
yield raw_f
else:
f = io.BufferedWriter(raw_f)
yield f
f.flush()
",point_e\util\ply_util.py
preprocess,,"def preprocess(data, channel):
if channel in COLORS:
return np.round(data * 255.0)
return data
",point_e\util\point_cloud.py
load,Load the point cloud from a .npz file.,"@classmethod
def load(cls, f: Union[str, BinaryIO]) ->'PointCloud':
""""""""""""
if isinstance(f, str):
with open(f, 'rb') as reader:
return cls.load(reader)
else:
obj = np.load(f)
keys = list(obj.keys())
return PointCloud(coords=obj['coords'], channels={k: obj[k] for k in
keys if k != 'coords'})
",point_e\util\point_cloud.py
save,Save the point cloud to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
""""""""""""
if isinstance(f, str):
with open(f, 'wb') as writer:
self.save(writer)
else:
np.savez(f, coords=self.coords, **self.channels)
",point_e\util\point_cloud.py
write_ply,,"def write_ply(self, raw_f: BinaryIO):
write_ply(raw_f, coords=self.coords, rgb=np.stack([self.channels[x] for
x in 'RGB'], axis=1) if all(x in self.channels for x in 'RGB') else
None)
",point_e\util\point_cloud.py
random_sample,"Sample a random subset of this PointCloud.
:param num_points: maximum number of points to sample.
:param subsample_kwargs: arguments to self.subsample().
:return: a reduced PointCloud, or self if num_points is not less than
the current number of points.","def random_sample(self, num_points: int, **subsample_kwargs) ->'PointCloud':
""""""""""""
if len(self.coords) <= num_points:
return self
indices = np.random.choice(len(self.coords), size=(num_points,),
replace=False)
return self.subsample(indices, **subsample_kwargs)
",point_e\util\point_cloud.py
farthest_point_sample,"Sample a subset of the point cloud that is evenly distributed in space.
First, a random point is selected. Then each successive point is chosen
such that it is furthest from the currently selected points.
The time complexity of this operation is O(NM), where N is the original
number of points and M is the reduced number. Therefore, performance
can be improved by randomly subsampling points with random_sample()
before running farthest_point_sample().
:param num_points: maximum number of points to sample.
:param init_idx: if specified, the first point to sample.
:param subsample_kwargs: arguments to self.subsample().
:return: a reduced PointCloud, or self if num_points is not less than
the current number of points.","def farthest_point_sample(self, num_points: int, init_idx: Optional[int]=
None, **subsample_kwargs) ->'PointCloud':
""""""""""""
if len(self.coords) <= num_points:
return self
init_idx = random.randrange(len(self.coords)
) if init_idx is None else init_idx
indices = np.zeros([num_points], dtype=np.int64)
indices[0] = init_idx
sq_norms = np.sum(self.coords ** 2, axis=-1)
def compute_dists(idx: int):
return sq_norms + sq_norms[idx] - 2 * (self.coords @ self.coords[idx])
cur_dists = compute_dists(init_idx)
for i in range(1, num_points):
idx = np.argmax(cur_dists)
indices[i] = idx
cur_dists = np.minimum(cur_dists, compute_dists(idx))
return self.subsample(indices, **subsample_kwargs)
",point_e\util\point_cloud.py
subsample,,"def subsample(self, indices: np.ndarray, average_neighbors: bool=False
) ->'PointCloud':
if not average_neighbors:
return PointCloud(coords=self.coords[indices], channels={k: v[
indices] for k, v in self.channels.items()})
new_coords = self.coords[indices]
neighbor_indices = PointCloud(coords=new_coords, channels={}
).nearest_points(self.coords)
neighbor_indices[indices] = np.arange(len(indices))
new_channels = {}
for k, v in self.channels.items():
v_sum = np.zeros_like(v[:len(indices)])
v_count = np.zeros_like(v[:len(indices)])
np.add.at(v_sum, neighbor_indices, v)
np.add.at(v_count, neighbor_indices, 1)
new_channels[k] = v_sum / v_count
return PointCloud(coords=new_coords, channels=new_channels)
",point_e\util\point_cloud.py
select_channels,,"def select_channels(self, channel_names: List[str]) ->np.ndarray:
data = np.stack([preprocess(self.channels[name], name) for name in
channel_names], axis=-1)
return data
",point_e\util\point_cloud.py
nearest_points,"For each point in another set of points, compute the point in this
pointcloud which is closest.
:param points: an [N x 3] array of points.
:param batch_size: the number of neighbor distances to compute at once.
Smaller values save memory, while larger values may
make the computation faster.
:return: an [N] array of indices into self.coords.","def nearest_points(self, points: np.ndarray, batch_size: int=16384
) ->np.ndarray:
""""""""""""
norms = np.sum(self.coords ** 2, axis=-1)
all_indices = []
for i in range(0, len(points), batch_size):
batch = points[i:i + batch_size]
dists = norms + np.sum(batch ** 2, axis=-1)[:, None] - 2 * (batch @
self.coords.T)
all_indices.append(np.argmin(dists, axis=-1))
return np.concatenate(all_indices, axis=0)
",point_e\util\point_cloud.py
combine,,"def combine(self, other: 'PointCloud') ->'PointCloud':
assert self.channels.keys() == other.channels.keys()
return PointCloud(coords=np.concatenate([self.coords, other.coords],
axis=0), channels={k: np.concatenate([v, other.channels[k]], axis=0
) for k, v in self.channels.items()})
",point_e\util\point_cloud.py