From 85ae985b6a232f3a3e2f7400243cec2ca0b5f8d1 Mon Sep 17 00:00:00 2001 From: Vladislav Veklenko <71467601+vladoossss@users.noreply.github.com> Date: Thu, 13 Oct 2022 01:44:01 +0200 Subject: [PATCH] update mask2segments and saving results (#9785) * update mask2segments and saving results * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update general.py Signed-off-by: Glenn Jocher * Update predict.py Signed-off-by: Glenn Jocher Signed-off-by: Glenn Jocher Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- utils/segment/general.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/utils/segment/general.py b/utils/segment/general.py index 43bdc46..b526333 100644 --- a/utils/segment/general.py +++ b/utils/segment/general.py @@ -126,9 +126,12 @@ def masks2segments(masks, strategy='largest'): segments = [] for x in masks.int().cpu().numpy().astype('uint8'): c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] - if strategy == 'concat': # concatenate all segments - c = np.concatenate([x.reshape(-1, 2) for x in c]) - elif strategy == 'largest': # select largest segment - c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2) + if c: + if strategy == 'concat': # concatenate all segments + c = np.concatenate([x.reshape(-1, 2) for x in c]) + elif strategy == 'largest': # select largest segment + c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2) + else: + c = np.zeros((0, 2)) # no segments found segments.append(c.astype('float32')) return segments