make thresholds backwards compatible

Author: ssmmoo1

examples/owl_predict.py

@@ -31,7 +31,7 @@
     parser = argparse.ArgumentParser()
     parser.add_argument("--image", type=str, default="../assets/owl_glove_small.jpg")
     parser.add_argument("--prompt", type=str, default="an owl, a glove")
-    parser.add_argument("--thresholds", type=str, default="0.1,0.1")
+    parser.add_argument("--threshold", type=str, default="0.1,0.1")
     parser.add_argument("--output", type=str, default="../data/owl_predict_out.jpg")
     parser.add_argument("--model", type=str, default="google/owlvit-base-patch32")
     parser.add_argument("--image_encoder_engine", type=str, default="../data/owl_image_encoder_patch32.engine")
@@ -43,7 +43,7 @@
     text = prompt.split(',')
     print(text)
 
-    thresholds = args.thresholds.strip("][()")
+    thresholds = args.threshold.strip("][()")
     thresholds = thresholds.split(',')
     thresholds = [float(x) for x in thresholds]
     print(thresholds)
@@ -62,7 +62,7 @@
         image=image, 
         text=text, 
         text_encodings=text_encodings,
-        thresholds=thresholds,
+        threshold=thresholds,
         pad_square=False
     )
 

examples/tree_predict.py

@@ -33,7 +33,7 @@
     parser.add_argument("--threshold", type=float, default=0.1)
     parser.add_argument("--output", type=str, default="../data/tree_predict_out.jpg")
     parser.add_argument("--model", type=str, default="google/owlvit-base-patch32")
-    parser.add_argument("--image_encoder_engine", type=str, default="../data/owlvit_image_encoder_patch32.engine")
+    parser.add_argument("--image_encoder_engine", type=str, default="../data/owl_image_encoder_patch32.engine")
     args = parser.parse_args()
 
     predictor = TreePredictor(

nanoowl/owl_predictor.py

@@ -274,9 +274,12 @@ def encode_rois(self, image: torch.Tensor, rois: torch.Tensor, pad_square: bool
     def decode(self, 
             image_output: OwlEncodeImageOutput, 
             text_output: OwlEncodeTextOutput,
-            thresholds: List[float],
+            threshold: Union[int, float, List[Union[int, float]]] = 0.1,
         ) -> OwlDecodeOutput:
 
+        if isinstance(threshold, (int, float)):
+            threshold = [threshold]
+
         num_input_images = image_output.image_class_embeds.shape[0]
 
         image_class_embeds = image_output.image_class_embeds
@@ -291,9 +294,9 @@ def decode(self,
         labels = scores_max.indices
         scores = scores_max.values
         masks = []
-        for i, threshold in enumerate(thresholds):
+        for i, thresh in enumerate(threshold):
             label_mask = labels == i
-            score_mask = scores > threshold 
+            score_mask = scores > thresh
             obj_mask = torch.logical_and(label_mask,score_mask)
             masks.append(obj_mask) 
 
@@ -455,7 +458,7 @@ def predict(self,
             image: PIL.Image, 
             text: List[str], 
             text_encodings: Optional[OwlEncodeTextOutput],
-            thresholds: List[float],
+            threshold: Union[int, float, List[Union[int, float]]] = 0.1,
             pad_square: bool = True,
 
         ) -> OwlDecodeOutput:
@@ -469,5 +472,5 @@ def predict(self,
 
         image_encodings = self.encode_rois(image_tensor, rois, pad_square=pad_square)
 
-        return self.decode(image_encodings, text_encodings, thresholds)
+        return self.decode(image_encodings, text_encodings, threshold)