From 31e25c8f99da64ecebc6591d481d65ef861c7cfe Mon Sep 17 00:00:00 2001 From: Gyeongbo Sim Date: Wed, 20 Jul 2022 22:00:07 +0900 Subject: [PATCH] fix: Use confidence threshold in demo.py --- demo/predictor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/demo/predictor.py b/demo/predictor.py index 189ec797..2d336bcc 100644 --- a/demo/predictor.py +++ b/demo/predictor.py @@ -35,6 +35,7 @@ def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False): self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu) else: self.predictor = DefaultPredictor(cfg) + self.confidence_threshold = cfg.confidence_threshold def run_on_image(self, image): """ @@ -61,7 +62,7 @@ def run_on_image(self, image): predictions["sem_seg"].argmax(dim=0).to(self.cpu_device) ) if "instances" in predictions: - instances = predictions["instances"].to(self.cpu_device) + instances = predictions["instances"][predictions["instances"].score >= self.confidence_threshold].to(self.cpu_device) vis_output = visualizer.draw_instance_predictions(predictions=instances) return predictions, vis_output