def to_stl(file, x, y, z):
# run octave
octave.addpath(".")
octave.stlwrite(f"../meshes/{file}.stl", x, y, z)
if __name__ == "__main__":
# model output (voxel volume) -> obj -> meshlab visualization
os.makedirs("../meshes", exist_ok=True)
args = get_args()
data = data_loader.FacesWith3DCoords(
images_dir=args.images_dir,
mats_dir=args.mats_dir,
landmarks_dir=args.lands_dir,
transform=args.transform,
)
model = StackedHourGlass(nChannels=224,
nStack=2,
nModules=2,
numReductions=4,
nOutputs=200)
model.cuda()
model.load_state_dict(torch.load(args.checkpoint))
model.eval()
print(f"[+] Loaded model {args.checkpoint}")
i = np.random.randint(len(data))
img, true_vol = data[i]
import torch
import torch.nn.functional as F
import torch.optim as optim
import data_loader
from model_hourglass import StackedHourGlass
from utils import get_args, to_cuda
args = get_args()
trainset = data_loader.FacesWith3DCoords(images_dir=args.images_dir,
mats_dir=args.mats_dir,
transform=args.transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=2,
shuffle=True,
num_workers=2)
model = StackedHourGlass(nChannels=224,
nStack=2,
nModules=2,
numReductions=4,
nOutputs=200)
model.cuda()
model.train()
print("Set-up model")
# criterion & optimizer
criterion = torch.nn.CrossEntropyLoss()
# optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9)