Spaces:
Sleeping
Sleeping
File size: 4,388 Bytes
757ed1c 5c52e2e 757ed1c bdfffc9 757ed1c 5c52e2e 757ed1c 5c52e2e 757ed1c 5c52e2e 757ed1c 5c52e2e 757ed1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import utils.binvox_rw as binvox_rw
import numpy as np
import plotly.graph_objects as go
from models.encoder import Encoder
from models.decoder import Decoder
from models.merger import Merger
from models.refiner import Refiner
from config import cfg
import torch
from datetime import datetime as dt
import utils.data_transforms
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
# device='cpu'
cfg.CONST.WEIGHTS='saved_model/Pix2Vox.pth'
def read_binvox(file) -> np.ndarray:
model = binvox_rw.read_as_3d_array(file)
return model.data.astype(np.uint8)
def voxel_to_plotly(voxels):
x, y, z = voxels.nonzero()
fig = go.Figure(data=[
go.Scatter3d(
x=x, y=y, z=z,
mode='markers',
marker=dict(size=3, color=z, colorscale='Viridis', opacity=0.7)
)
])
fig.update_layout(scene=dict(aspectmode='data'))
return fig
# when gpu is not available
def remove_module_prefix(state_dict):
new_state_dict = {}
for k, v in state_dict.items():
if k.startswith('module.'):
new_state_dict[k[7:]] = v
else:
new_state_dict[k] = v
return new_state_dict
IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W
CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W
test_transforms = utils.data_transforms.Compose([
utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE),
utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE),
utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD),
utils.data_transforms.ToTensor(),
])
def predict_voxel_from_images(rendering_images):
transformed_images = test_transforms(rendering_images)
encoder = Encoder(cfg)
decoder = Decoder(cfg)
refiner = Refiner(cfg)
merger = Merger(cfg)
print('[INFO] %s Loading weights from %s ...' % (dt.now(), cfg.CONST.WEIGHTS))
checkpoint = torch.load(cfg.CONST.WEIGHTS, weights_only=False, map_location=device)
if torch.cuda.is_available():
encoder = torch.nn.DataParallel(encoder).cuda()
decoder = torch.nn.DataParallel(decoder).cuda()
refiner = torch.nn.DataParallel(refiner).cuda()
merger = torch.nn.DataParallel(merger).cuda()
encoder_state_dict=checkpoint['encoder_state_dict']
decoder_state_dict=checkpoint['decoder_state_dict']
merger_state_dict=checkpoint['merger_state_dict']
refiner_state_dict = checkpoint['refiner_state_dict']
else:
encoder_state_dict = remove_module_prefix(checkpoint['encoder_state_dict'])
decoder_state_dict = remove_module_prefix(checkpoint['decoder_state_dict'])
merger_state_dict = remove_module_prefix(checkpoint['merger_state_dict'])
refiner_state_dict = remove_module_prefix(checkpoint['refiner_state_dict'])
epoch_idx = checkpoint['epoch_idx']
encoder.load_state_dict(encoder_state_dict)
decoder.load_state_dict(decoder_state_dict)
if cfg.NETWORK.USE_REFINER:
refiner.load_state_dict(refiner_state_dict)
if cfg.NETWORK.USE_MERGER:
merger.load_state_dict(merger_state_dict)
encoder.eval()
decoder.eval()
merger.eval()
refiner.eval()
with torch.no_grad():
transformed_images = transformed_images.unsqueeze(0) #adding the batch_dim
transformed_images = transformed_images.to(device)
# print(rendering_images.shape)
image_features = encoder(transformed_images)
print(image_features.shape)
raw_features, generated_volume = decoder(image_features)
print(generated_volume.shape)
if cfg.NETWORK.USE_MERGER:
generated_volume = merger(raw_features, generated_volume)
else:
generated_volume = torch.mean(generated_volume, dim=1)
# encoder_loss = bce_loss(generated_volume, ground_truth_volume) * 10
if cfg.NETWORK.USE_REFINER:
generated_volume = refiner(generated_volume)
# refiner_loss = bce_loss(generated_volume, ground_truth_volume) * 10
else:
# refiner_loss = encoder_loss
pass
generated_volume=generated_volume.squeeze(0)
gv = generated_volume.cpu().numpy()
gv = (gv >= 0.5).astype(np.uint8)
torch.cuda.empty_cache()
return gv
|