natanielruiz
2017-09-14 ec99c6649af6bdbd3c836f20cdc81170e7045cc8
code/test_on_video.py
@@ -60,7 +60,7 @@
    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)
@@ -141,7 +141,7 @@
        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)