[ PyTorch / torchvision ] draw_bounding_boxes() 사용하기

---

[ PyTorch / torchvision ] draw_bounding_boxes() 사용하기

 

https://pytorch.org/

---

[ 이전 글 ]

[ Pytorch ] 파이토치 텐서 합치는 방법 : cat(), stack() ( + dim의 의미와, 병합 방식의 차이)

[ Pytorch ] nn.BCELoss, nn.BCEWithLogitsLoss, nn.CrossEntropyLoss, nn.NLLLoss 총정리

[ PyTorch / torchvision ] make_grid() 사용하기

---

torchivision.utils.draw_bounding_box (링크)

이미지에 Bounding Box를 그려주는 함수입니다.

Input 이미지는 0~255 사이의 값(uint8)을 가져야 합니다.

parameters 

• tensor (Tensor) : Tensor (Channel, Height, Width)
• boxes (Tensor) : Tensor (N, 4). N : bounding box 개수, 4 : (xmin, ymin, xmax, ymax) format. 값은 이미지의 절대 좌표 (0 <= xmin < xmax < Width, 0 <= ymin < ymax < Height)
• labels (List[str]) : Bounding Box의 라벨을 리스트로 설정
• colors (color or list of colors, optional) : Bounding Box의 색을 담은 리스트 또는 전체 바운딩 박스 색을 통일. PIL string(e.g. "red" or "#FF00FF", RGB tuple(e.g. (240, 10, 157)) 등으로 표현. (Default : 박스마다 랜덤 컬러)
• fill (bool) : True 일 경우, 특정한 색으로 바운딩 박스를 채움 (Default : False)
• width (int) : Bounding Box의 선 두께
• font (str) : TrueType font 를 포함한 파일 이름. 없다면, 다른 폰트에서 찾아서 아무거나 적용.
• font_size (int) : 글자 크기

 

returns

• tensor : 바운딩 박스가 표현된 uint8 이미지 텐서

---

예시 

• 1 ) 시각화 코드
• 2 ) draw_bounding_boxes 적용
• 3 ) draw_bounding_boxes 적용 : color 설정
• 4 ) draw_bounding_boxes 적용 : fill 설정
• 5 ) 여러 사진에 Bounding Box 그리기

사용된 사진 : dog1.jpg, dog2.jpg

아래의 사진을 다운 받으셔서 assets 이라는 파일 아래에 저장해주세요.

ex) assets/dog1.jpg

from torch/vision

1 ) 시각화 코드

draw_bounding_boxes 의 결과를 확인하기 위한 시각화 코드입니다.

# code from https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html

import torch
import numpy as np
import matplotlib.pyplot as plt

import torchvision.transforms.functional as F

plt.rcParams["savefig.bbox"] = 'tight'

def show(imgs):
    if not isinstance(imgs, list): # 하나의 이미지일때
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False) # 총 사진의 개수만큼 plot
    for i, img in enumerate(imgs):
        img = img.detach() # 학습 그래프에서 제외
        img = F.to_pil_image(img) # torch.tensor 에서 pil 이미지로 변환
        axs[0, i].imshow(np.asarray(img)) # numpy 배열로 변경후, 가로로 이미지를 나열
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

위의 코드를 이해하기 위해서는 아래의 글을 참고하시면 좋습니다.

 

[ Python / PIL ] PIL 이미지와 Torch.Tensor 변환 (ToTensor, ToPILImage)

[ Python / PIL ] PIL 이미지와 Torch.Tensor 변환 (ToTensor, ToPILImage)

 

2 ) draw_bounding_boxes 적용

Bounding Box를 2개 그리겠습니다.

boxes : (2 , (xmin, ymin, xmax, ymax) format)

width : 바운딩 박스의 두께는 5포인트로 설정

# code from https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html

from torchvision.io import read_image
from pathlib import Path
from torchvision.utils import draw_bounding_boxes

# read_image : 이미지 파일을 읽어옵니다.
dog1_int = read_image(str(Path('assets') / 'dog1.jpg')) # assets/dog1.jpg

boxes = torch.tensor([[50, 50, 100, 200], [210, 150, 350, 430]], dtype=torch.float)
result = draw_bounding_boxes(dog1_int, boxes, width=5)
show(result)

결과

• 따로 color 을 지정해주지 않았기 때문에 랜덤으로 검은색([50,50,100,200])과 초록색([210,150,350,430])이 사용되었습니다. 

 

3 ) draw_bounding_boxes 적용 : color 설정

이번에는 color 를 추가하겠습니다.

PIL string(e.g. "red" or "#FF00FF" 또는 RGB tuple(e.g. (240, 10, 157)) 등으로 표현가능 합니다.

 

# code from https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html

from torchvision.io import read_image
from pathlib import Path
from torchvision.utils import draw_bounding_boxes

# read_image : 이미지 파일을 읽어옵니다.
dog1_int = read_image(str(Path('assets') / 'dog1.jpg')) # assets/dog1.jpg

boxes = torch.tensor([[50, 50, 100, 200], [210, 150, 350, 430]], dtype=torch.float)
colors = ["blue", "yellow"] # 색 지정
result = draw_bounding_boxes(dog1_int, boxes, colors=colors, width=5)
show(result)

결과

• 푸른색([50,50,100,200])과 노란색([210,150,350,430])이 사용되었습니다. 

 

4 ) draw_bounding_boxes 적용 : fill 설정

이번에는 fill 을 추가하겠습니다.

True 일 경우, 특정한 색으로 바운딩 박스를 채움 (Default : False)

# code from https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html
from torchvision.utils import make_grid
from torchvision.io import read_image
from pathlib import Path
from torchvision.utils import draw_bounding_boxes

# read_image : 이미지 파일을 읽어옵니다.
dog1_int = read_image(str(Path('assets') / 'dog1.jpg')) # assets/dog1.jpg

boxes = torch.tensor([[50, 50, 100, 200], [210, 150, 350, 430]], dtype=torch.float)
colors = ["blue", "yellow"] # 색 지정
result = draw_bounding_boxes(dog1_int, boxes, colors=colors, fill=True, width=5)
show(result)

결과

• 푸른색([50,50,100,200])과 노란색([210,150,350,430])으로 Bounding 박스가 채워졌습니다.

 

 

5 ) 여러 사진에 Bounding Box 그리기

보통 Deep Learning Model 의 경우 여러개의 이미지에 대한 bounding box 가 동시에 제시되기도 합니다.

이럴 경우엔 어떻게 Bounding Box 함수를 적용해야 할까요?

 

사전학습된 모델을 불러와서 직접 적용해 봅시다.

# code from https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html

from torchvision.io import read_image
from pathlib import Path

dog1_int = read_image(str(Path('assets') / 'dog1.jpg'))
dog2_int = read_image(str(Path('assets') / 'dog2.jpg'))
dog_list = [dog1_int, dog2_int]

from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights
# torchivision version : 0.13

# 사전학습된 객체 검출 모델 사용
weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
transforms = weights.transforms()

# 이미지 전처리
images = [transforms(d) for d in dog_list]

# 사전학습된 객체 검출 모델 사용
model = fasterrcnn_resnet50_fpn(weights=weights, progress=False)
model = model.eval()

outputs = model(images)
print(outputs)

결과

• 각 이미지에 대한 Bounding Box들과 그것들의 label, score 값을 구했습니다.

 

이때 bounding box score 가 0.8 이상인 것만 그려 봅시다.

• zip() 이라는 함수를 사용하여서 각 이미지마다 draw_bounding_boxes 를 적용할 수 있습니다.

from torchvision.utils import draw_bounding_boxes

score_threshold = .8
dogs_with_boxes = [
    draw_bounding_boxes(dog_int, boxes=output['boxes'][output['scores'] > score_threshold], width=4)
    for dog_int, output in zip(dog_list, outputs)
]
show(dogs_with_boxes)

결과

• 강아지 객체에만 검은색 Bounding Box 가 그려진 것을 볼 수 있습니다.

---

[ 다음 글 ]

[ PyTorch / torchvision ] draw_segmentation_masks() 사용하기

[ PyTorch / torchvision ] draw_segmentation_masks() 사용하기

---

[ 추가 내용 ] PIL 을 이용한 Bounding Box 그리기

[ Python / PIL ] PIL Image 에 바운딩 박스 그리기 (ImageDraw.Draw(), draw.rectangle(), grayscale)

[ Python / PIL ] PIL Image 에 바운딩 박스 그리기 (ImageDraw.Draw(), draw.rectangle(), grayscale)

---

[ 참고 ]

https://pytorch.org/vision/stable/auto_examples/plot_visualization_utils.html

---