def load_sdf_net(filename=None, return_latent_codes=False):
from model.sdf_net import SDFNet, LATENT_CODES_FILENAME
sdf_net = SDFNet()
if filename is not None:
sdf_net.filename = filename
sdf_net.load()
sdf_net.eval()
if return_latent_codes:
latent_codes = torch.load(LATENT_CODES_FILENAME).to(device)
latent_codes.requires_grad = False
return sdf_net, latent_codes
else:
return sdf_net
voxels_current,
level=0,
spacing=(size / voxel_resolution, size / voxel_resolution,
size / voxel_resolution))
vertices -= size / 2
mesh = trimesh.Trimesh(vertices=vertices,
faces=faces,
vertex_normals=normals)
point_cloud = mesh.sample(point_cloud_size)
rescale_point_cloud(point_cloud, method=rescale)
result[i, :, :] = point_cloud
return result
if 'sample' in sys.argv:
sdf_net = SDFNet()
sdf_net.filename = 'hybrid_gan_generator.to'
sdf_net.load()
sdf_net.eval()
clouds = sample_point_clouds(sdf_net, 1000, 2048, voxel_resolution=32)
np.save('data/generated_point_cloud_sample.npy', clouds)
if 'checkpoints' in sys.argv:
import glob
from tqdm import tqdm
torch.manual_seed(1234)
files = glob.glob(
'models/checkpoints/hybrid_progressive_gan_generator_2-epoch-*.to',
recursive=True)
latent_codes = standard_normal_distribution.sample(
mesh = trimesh.load(MODEL_PATH)
points, sdf = sample_sdf_near_surface(mesh)
save_images = 'save' in sys.argv
if save_images:
viewer = MeshRenderer(start_thread=False, size=1080)
ensure_directory('images')
else:
viewer = MeshRenderer()
points = torch.tensor(points, dtype=torch.float32, device=device)
sdf = torch.tensor(sdf, dtype=torch.float32, device=device)
sdf.clamp_(-0.1, 0.1)
sdf_net = SDFNet(latent_code_size=LATENT_CODE_SIZE).to(device)
optimizer = torch.optim.Adam(sdf_net.parameters(), lr=1e-5)
BATCH_SIZE = 20000
latent_code = torch.zeros((BATCH_SIZE, LATENT_CODE_SIZE), device=device)
indices = torch.zeros(BATCH_SIZE, dtype=torch.int64, device=device)
positive_indices = (sdf > 0).nonzero().squeeze().cpu().numpy()
negative_indices = (sdf < 0).nonzero().squeeze().cpu().numpy()
step = 0
error_targets = np.logspace(np.log10(0.02), np.log10(0.0005), num=500)
image_index = 0
while True:
try:
from datasets import VoxelDataset
from torch.utils.data import DataLoader
LEARN_RATE = 0.00001
BATCH_SIZE = 8
CRITIC_UPDATES_PER_GENERATOR_UPDATE = 5
CRITIC_WEIGHT_LIMIT = 0.01
dataset = VoxelDataset.glob('data/chairs/voxels_32/**.npy')
dataset.rescale_sdf = False
data_loader = DataLoader(dataset,
shuffle=True,
batch_size=BATCH_SIZE,
num_workers=8)
generator = SDFNet()
generator.filename = 'hybrid_wgan_generator.to'
critic = Discriminator()
critic.filename = 'hybrid_wgan_critic.to'
critic.use_sigmoid = False
if "continue" in sys.argv:
generator.load()
critic.load()
LOG_FILE_NAME = "plots/hybrid_wgan_training.csv"
first_epoch = 0
if 'continue' in sys.argv:
log_file_contents = open(LOG_FILE_NAME, 'r').readlines()
first_epoch = len(log_file_contents)
VOXEL_RESOLUTION = RESOLUTIONS[ITERATION]
dataset = VoxelDataset.from_split(
'data/chairs/voxels_{:d}/{{:s}}.npy'.format(VOXEL_RESOLUTION),
'data/chairs/train.txt')
data_loader = DataLoader(dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4)
def get_generator_filename(iteration):
return 'hybrid_progressive_gan_generator_{:d}.to'.format(iteration)
generator = SDFNet(device='cpu')
discriminator = Discriminator()
if not CONTINUE and ITERATION > 0:
generator.filename = get_generator_filename(ITERATION - 1)
generator.load()
discriminator.set_iteration(ITERATION - 1)
discriminator.load()
discriminator.set_iteration(ITERATION)
generator.filename = get_generator_filename(ITERATION)
if CONTINUE:
generator.load()
discriminator.load()
if torch.cuda.device_count() > 1:
print("Using dataparallel with {:d} GPUs.".format(
torch.cuda.device_count()))