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 이상인 경우만 카운트 합니다.

 

반응형