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)