FCN confusion matrix visualization fix

Fix LINT check
NVIDIA · Mar 27, 2017 · 730c4a7 · 730c4a7
1 parent 5dccb29
commit 730c4a7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/digits/model/images/classification/views.py b/digits/model/images/classification/views.py
@@ -529,6 +529,11 @@ def classify_many():
                 'Unable to classify any image from the file')
 
         scores = last_output_data
+        # force correct 2D shape squeezing scores
+        for i in reversed(range(2, len(scores.shape))):
+            if scores.shape[i] == 1:
+                scores = np.squeeze(scores, axis=(i,))
+
         # take top 5
         indices = (-scores).argsort()[:, :5]