From 2df9160c6daa84830973a12a41f9109bc25e498d Mon Sep 17 00:00:00 2001 From: Roshan Rao Date: Thu, 22 Aug 2019 14:47:29 -0700 Subject: [PATCH] [fix] Fix detected objects in extract_features_vmb (#159) * [fix] Correctly ignore the background class * [fix] Output the correct number of objects --- pythia/scripts/features/extract_features_vmb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pythia/scripts/features/extract_features_vmb.py b/pythia/scripts/features/extract_features_vmb.py index c18566177..77988bb62 100644 --- 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( {