| | |
| | | print 'Loading data.' |
| | | |
| | | transformations = transforms.Compose([transforms.Scale(224), |
| | | transforms.RandomCrop(224), transforms.ToTensor(), |
| | | transforms.CenterCrop(224), transforms.ToTensor(), |
| | | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) |
| | | |
| | | model.cuda(gpu) |
| | |
| | | img_shape = img.size() |
| | | img = img.view(1, img_shape[0], img_shape[1], img_shape[2]) |
| | | img = Variable(img).cuda(gpu) |
| | | yaw, pitch, roll = model(img) |
| | | yaw, pitch, roll, angles = model(img) |
| | | |
| | | yaw_predicted = F.softmax(yaw) |
| | | pitch_predicted = F.softmax(pitch) |