Mask R-CNN으로 인스턴스 분할

PyTorch로 배우는 이미지 딥러닝

Michal Oleszak

Machine Learning Engineer

Faster R-CNN

Faster R-CNN 아키텍처 다이어그램

Mask R-CNN

Mask R-CNN 아키텍처 다이어그램

PyTorch의 사전학습 Mask R-CNN

from torchvision.models.detection import \
maskrcnn_resnet50_fpn


model = maskrcnn_resnet50_fpn(pretrained=True)
model.eval()


image = Image.open("cat_and_laptop.jpg")
transform = transforms.Compose([
    transforms.ToTensor()
])
image_tensor = transform(image).unsqueeze(0)


with torch.no_grad():
    prediction = model(image_tensor)

Mask R-CNN 모델 임포트
사전학습 모델 불러오기
테스트 이미지를 불러와 텐서로 변환

노트북 옆에 앉아 있는 고양이 사진

이미지 텐서를 모델에 전달

모델 출력

레이블

prediction[0]["labels"]

tensor([
    17, 73, 76, 73, 67, 42, 63, 84,73, 65, 
    17, 73, 73, 73, 84, 72, 76, 76,17, 15
])

클래스 이름

print(class_names[17], class_names[73])

cat laptop

클래스 확률

prediction[0]["scores"]

tensor([
    0.9981, 0.9672, 0.9061, 0.6893, 0.3729, 
    ..., 
    0.0745, 0.0705, 0.0623, 0.0610, 0.0508
])

마스크

prediction[0]["masks"]

tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],
          ...]]]])

소프트 마스크

고유 마스크 값

prediction[0]["masks"].unique()

tensor([0.0000e+00, 5.9713e-08, ..., 
        9.9989e-01, 9.9990e-01])

Mask R-CNN 마스크:
- 값은 0~1 사이
- 각 픽셀이 객체에 속할 확신도를 나타냄
- 이진 마스크보다 정보가 풍부함
- 필요 시 임계값으로 이진화 가능

소프트 마스크 표시하기

masks = prediction[0]["masks"]
labels = prediction[0]["labels"]


for i in range(2):
    plt.imshow(image)

    plt.imshow(
        masks[i, 0],
        cmap="jet",
        alpha=0.5,
    )

    plt.title(
        f"Object: {class_names[labels[i]]}"
    )
    plt.show()

예측에서 마스크와 레이블 추출
상위 두 객체를 순회하며 원본 이미지를 표시
각 객체에 대해 반투명 마스크를 겹쳐 그림
제목 추가 후 표시

소프트 마스크 표시하기

이미지 위에 겹쳐진 인스턴스 분할 마스크

연습해 봅시다!

PyTorch로 배우는 이미지 딥러닝