JeffLiang commited on
Commit
304ea95
1 Parent(s): 85fdfea

fix bug w/ only one class inputs

Browse files
Files changed (1) hide show
  1. open_vocab_seg/utils/predictor.py +2 -0
open_vocab_seg/utils/predictor.py CHANGED
@@ -204,6 +204,8 @@ class SAMVisualizationDemo(object):
204
  if self.granularity < 1:
205
  thr_scores = max_scores * self.granularity
206
  select_mask = []
 
 
207
  for i, thr in enumerate(thr_scores):
208
  cls_pred = class_preds[:,i]
209
  locs = torch.where(cls_pred > thr)
 
204
  if self.granularity < 1:
205
  thr_scores = max_scores * self.granularity
206
  select_mask = []
207
+ if len(class_names) == 2 and class_names[-1] == 'others':
208
+ thr_scores = thr_scores[:-1]
209
  for i, thr in enumerate(thr_scores):
210
  cls_pred = class_preds[:,i]
211
  locs = torch.where(cls_pred > thr)