Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from pytorch3d.structures import Pointclouds | |
| def set_requires_grad(module: nn.Module, requires_grad: bool): | |
| for p in module.parameters(): | |
| p.requires_grad_(requires_grad) | |
| def compute_distance_transform(mask: torch.Tensor): | |
| """ | |
| Parameters | |
| ---------- | |
| mask (B, 1, H, W) or (B, 2, H, W) true for foreground | |
| Returns | |
| ------- | |
| the vector to the closest foreground pixel, zero if inside mask | |
| """ | |
| C = mask.shape[1] | |
| assert C in [1, 2], f'invalid mask shape {mask.shape} found!' | |
| image_size = mask.shape[-1] | |
| dts = [] | |
| for i in range(C): | |
| distance_transform = torch.stack([ | |
| torch.from_numpy(cv2.distanceTransform( | |
| (1 - m), distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_3 | |
| ) / (image_size / 2)) | |
| for m in mask[:, i:i+1].squeeze(1).detach().cpu().numpy().astype(np.uint8) | |
| ]).unsqueeze(1).clip(0, 1).to(mask.device) | |
| dts.append(distance_transform) | |
| return torch.cat(dts, 1) | |
| def default(x, d): | |
| return d if x is None else x | |
| def get_num_points(x: Pointclouds, /): | |
| return x.points_padded().shape[1] | |
| def get_custom_betas(beta_start: float, beta_end: float, warmup_frac: float = 0.3, num_train_timesteps: int = 1000): | |
| """Custom beta schedule""" | |
| betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32) | |
| warmup_frac = 0.3 | |
| warmup_time = int(num_train_timesteps * warmup_frac) | |
| warmup_steps = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) | |
| warmup_time = min(warmup_time, num_train_timesteps) | |
| betas[:warmup_time] = warmup_steps[:warmup_time] | |
| return betas | |