Python/PyTorch 공부
segmentation mask 덩어리 갯수 확인하기
AI 꿈나무
2022. 8. 21. 20:40
반응형
위 그림은 mask가 두 덩어리로 이루어 있습니다.
이 케이스를 파악하는 코드를 짜보았는뎅 공유합니다.
def count_contours(masks, threshold=1000):
if type(masks[0]) == torch.Tensor:
masks = [mask.cpu().numpy() for mask in masks]
counts = []
for mask in masks:
count = 0
contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
for contour in contours:
area = cv2.contourArea(contour)
if area >= threshold:
count += 1
counts.append(count)
return np.mean(counts)
cv2.findContours 함수를 사용하여 contour를 얻어옵니다.
contourArea 함수를 이용하여 contour의 영역 크기를 계산한 후, threshold 이상인 경우만 카운트 합니다.
반응형