faster process_mask(upsample=True) for predict

zldrobit · Aug 26, 2022 · 6bec10e · 6bec10e
1 parent f21e349
commit 6bec10e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/segment/predict.py b/segment/predict.py
@@ -44,7 +44,7 @@
 from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                            increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
 from utils.plots import Annotator, colors, save_one_box
-from utils.segment.general import process_mask_upsample, scale_masks
+from utils.segment.general import process_mask, scale_masks
 from utils.segment.plots import plot_masks
 from utils.torch_utils import select_device, smart_inference_mode
 
@@ -149,7 +149,7 @@ def run(
             annotator = Annotator(im0, line_width=line_thickness, example=str(names))
             if len(det):
                 # Mask additions ---------------------------------------------------------------------------------------
-                masks = process_mask_upsample(proto[i], det[:, 6:], det[:, :4], im.shape[2:])  # HWC
+                masks = process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True)  # HWC
                 masks = masks.permute(2, 0, 1).contiguous()  # CHW
                 # Mask additions ---------------------------------------------------------------------------------------