wip: add interactive seg model

This commit is contained in:
Qing
2022-11-27 21:25:27 +08:00
parent af87cca643
commit 023306ae40
20 changed files with 820 additions and 46 deletions

View File

@@ -180,3 +180,29 @@ def boxes_from_mask(mask: np.ndarray) -> List[np.ndarray]:
boxes.append(box)
return boxes
def only_keep_largest_contour(mask: np.ndarray) -> List[np.ndarray]:
"""
Args:
mask: (h, w) 0~255
Returns:
"""
_, thresh = cv2.threshold(mask, 127, 255, 0)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
max_area = 0
max_index = -1
for i, cnt in enumerate(contours):
area = cv2.contourArea(cnt)
if area > max_area:
max_area = area
max_index = i
if max_index != -1:
new_mask = np.zeros_like(mask)
return cv2.drawContours(new_mask, contours, max_index, 255, -1)
else:
return mask