From 2eb13d63b15a8ac908d6fa324c7f3d19141ca570 Mon Sep 17 00:00:00 2001
From: natanielruiz <nataniel777@hotmail.com>
Date: 星期六, 12 八月 2017 08:57:15 +0800
Subject: [PATCH] Temperature softmax and 10 shape PCA regression.

---
 code/test_resnet_bins.py |   99 ++++++++++++++++++++++++++++++++++++++-----------
 1 files changed, 77 insertions(+), 22 deletions(-)

diff --git a/code/test_resnet_bins.py b/code/test_resnet_bins.py
index 34fc8f5..4b1a655 100644
--- a/code/test_resnet_bins.py
+++ b/code/test_resnet_bins.py
@@ -6,6 +6,7 @@
 from torchvision import transforms
 import torch.backends.cudnn as cudnn
 import torchvision
+import torch.nn.functional as F
 
 import cv2
 import matplotlib.pyplot as plt
@@ -13,7 +14,7 @@
 import os
 import argparse
 
-from datasets import AFLW2000
+import datasets
 import hopenet
 import utils
 
@@ -30,6 +31,8 @@
           default='', type=str)
     parser.add_argument('--batch_size', dest='batch_size', help='Batch size.',
           default=1, type=int)
+    parser.add_argument('--save_viz', dest='save_viz', help='Save images with pose cube.',
+          default=False, type=bool)
 
     args = parser.parse_args()
 
@@ -39,14 +42,15 @@
     args = parse_args()
 
     cudnn.enabled = True
-    batch_size = 1
     gpu = args.gpu_id
     snapshot_path = os.path.join('output/snapshots', args.snapshot + '.pkl')
 
-    model = torchvision.models.resnet18()
-    # Parameters of newly constructed modules have requires_grad=True by default
-    num_ftrs = model.fc.in_features
-    model.fc = nn.Linear(num_ftrs, 3)
+    # ResNet101 with 3 outputs.
+    # model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 23, 3], 66)
+    # ResNet50
+    model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66)
+    # ResNet18
+    # model = hopenet.Hopenet(torchvision.models.resnet.BasicBlock, [2, 2, 2, 2], 66)
 
     print 'Loading snapshot.'
     # Load snapshot
@@ -55,12 +59,13 @@
 
     print 'Loading data.'
 
-    transformations = transforms.Compose([transforms.Scale(224),transforms.RandomCrop(224), transforms.ToTensor()])
+    transformations = transforms.Compose([transforms.Scale(224),
+    transforms.RandomCrop(224), transforms.ToTensor()])
 
-    pose_dataset = AFLW2000(args.data_dir, args.filename_list,
+    pose_dataset = datasets.AFLW2000_binned(args.data_dir, args.filename_list,
                                 transformations)
     test_loader = torch.utils.data.DataLoader(dataset=pose_dataset,
-                                               batch_size=batch_size,
+                                               batch_size=args.batch_size,
                                                num_workers=2)
 
     model.cuda(gpu)
@@ -69,22 +74,72 @@
 
     # Test the Model
     model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
-    error = .0
     total = 0
+    n_margins = 20
+    yaw_correct = np.zeros(n_margins)
+    pitch_correct = np.zeros(n_margins)
+    roll_correct = np.zeros(n_margins)
+
+    idx_tensor = [idx for idx in xrange(66)]
+    idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu)
+
+    yaw_error = .0
+    pitch_error = .0
+    roll_error = .0
+
+    l1loss = torch.nn.L1Loss(size_average=False)
+
     for i, (images, labels, name) in enumerate(test_loader):
         images = Variable(images).cuda(gpu)
-        labels = Variable(labels).cuda(gpu)
-        outputs = model(images)
-        _, predicted = torch.max(outputs.data, 1)
         total += labels.size(0)
-        # TODO: There are more efficient ways.
-        for idx in xrange(len(outputs)):
-            # if abs(outputs[idx].data[1] - labels[idx].data[1]) * 180 / np.pi > 30:
-            print name
-            print abs(outputs[idx].data - labels[idx].data) * 180 / np.pi, 180 * outputs[idx].data / np.pi, labels[idx].data * 180 / np.pi
-            # error += utils.mse_loss(outputs[idx], labels[idx])
-            error += abs(outputs[idx].data - labels[idx].data) * 180 / np.pi
+        label_yaw = labels[:,0].float()
+        label_pitch = labels[:,1].float()
+        label_roll = labels[:,2].float()
 
+        yaw, pitch, roll = model(images)
 
-    print('Test MSE error of the model on the ' + str(total) +
-    ' test images: %.4f' % (error / total))
+        # Binned predictions
+        _, yaw_bpred = torch.max(yaw.data, 1)
+        _, pitch_bpred = torch.max(pitch.data, 1)
+        _, roll_bpred = torch.max(roll.data, 1)
+
+        # Continuous predictions
+        yaw_predicted = utils.softmax_temperature(yaw.data, 1)
+        pitch_predicted = utils.softmax_temperature(pitch.data, 1)
+        roll_predicted = utils.softmax_temperature(roll.data, 1)
+
+        yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu()
+        pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu()
+        roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu()
+
+        # Mean absolute error
+        yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw) * 3)
+        pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch) * 3)
+        roll_error += torch.sum(torch.abs(roll_predicted - label_roll) * 3)
+
+        # Binned Accuracy
+        # for er in xrange(n_margins):
+        #     yaw_bpred[er] += (label_yaw[0] in range(yaw_bpred[0,0] - er, yaw_bpred[0,0] + er + 1))
+        #     pitch_bpred[er] += (label_pitch[0] in range(pitch_bpred[0,0] - er, pitch_bpred[0,0] + er + 1))
+        #     roll_bpred[er] += (label_roll[0] in range(roll_bpred[0,0] - er, roll_bpred[0,0] + er + 1))
+
+        # print label_yaw[0], yaw_bpred[0,0]
+
+        # Save images with pose cube.
+        # TODO: fix for larger batch size
+        if args.save_viz:
+            name = name[0]
+            cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg'))
+            #print os.path.join('output/images', name + '.jpg')
+            #print label_yaw[0] * 3 - 99, label_pitch[0] * 3 - 99, label_roll[0] * 3 - 99
+            #print yaw_predicted * 3 - 99, pitch_predicted * 3 - 99, roll_predicted * 3 - 99
+            utils.plot_pose_cube(cv2_img, yaw_predicted[0] * 3 - 99, pitch_predicted[0] * 3 - 99, roll_predicted[0] * 3 - 99)
+            cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img)
+
+    print('Test error in degrees of the model on the ' + str(total) +
+    ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f' % (yaw_error / total,
+    pitch_error / total, roll_error / total))
+
+    # Binned accuracy
+    # for idx in xrange(len(yaw_correct)):
+    #     print yaw_correct[idx] / total, pitch_correct[idx] / total, roll_correct[idx] / total

--
Gitblit v1.8.0