[fix] Fix detected objects in extract_features_vmb (#159)

* [fix] Correctly ignore the background class * [fix] Output the correct number of objects
facebookresearch · Aug 22, 2019 · 2df9160 · 2df9160
1 parent 12f67cd
commit 2df9160
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pythia/scripts/features/extract_features_vmb.py b/pythia/scripts/features/extract_features_vmb.py
@@ -162,7 +162,7 @@ def _process_feature_extraction(
             feat_list.append(feats[i][keep_boxes])
             bbox = output[0]["proposals"][i][keep_boxes].bbox / im_scales[i]
             # Predict the class label using the scores
-            objects = torch.argmax(scores[keep_boxes][start_index:], dim=1)
+            objects = torch.argmax(scores[keep_boxes][:, start_index:], dim=1)
 
             info_list.append(
                 {