[pytorch]Box operator로 Bounding Box 쉽게 다루기

torchvision.ops

torchvision.ops 에서는 computer vision 관련 각종 operator 들을 제공하는데, 그중에서 Bounding Box를 다루는 Box Operators를 소개한다. Detection 관련 코드를 만지다보면 bounding box를 다룰 일이 많은데, torchvision.ops 에서는 bactch 단위로 연산이 가능한 각종 기능들을 제공한다.

 

* box 면적, box간의 iou, box format 변경, clip 등 수행 가능

 

 

테스트

• box 3개를 정의하고 앞서 살펴본 box operator 기능들을 테스트

import numpy as np
import cv2
import torch
from torchvision.ops import *

# boxes shape : (N, 4)
# box format : x1, y1, x2, y2
boxes = torch.tensor([[50,50,150,110],      
                      [70,30,200,80],
                      [200,300,250,400]])

black_canvas = np.zeros((500,500,3))

cv2.rectangle(black_canvas, (boxes[0][0].item(), boxes[0][1].item()), (boxes[0][2].item(), boxes[0][3].item()), (0,0,255),5)
cv2.rectangle(black_canvas, (boxes[1][0].item(), boxes[1][1].item()), (boxes[1][2].item(), boxes[1][3].item()), (0,255,0),5)
cv2.rectangle(black_canvas, (boxes[2][0].item(), boxes[2][1].item()), (boxes[2][2].item(), boxes[2][3].item()), (255,0,0),5)

print('- boxes.shape :', boxes.shape)
print('- box_area :',box_area(boxes))
print('- box_convert (xyxy -> xywh) :\n',box_convert(boxes, in_fmt='xyxy', out_fmt='xywh'))
print('- box_convert (xyxy -> cxcywh) :\n',box_convert(boxes, in_fmt='xyxy', out_fmt='cxcywh'))
print('- box_iou (boxes[0], boxes[1]) :', box_iou(boxes[0].unsqueeze(0), boxes[1].unsqueeze(0)))
print('- box_iou (boxes[0], boxes[2]) :', box_iou(boxes[0].unsqueeze(0), boxes[2].unsqueeze(0)))
print('- generalized_box_iou (boxes[0], boxes[1]) :', generalized_box_iou(boxes[0].unsqueeze(0), boxes[1].unsqueeze(0)))
print('- generalized_box_iou (boxes[0], boxes[2]) :', generalized_box_iou(boxes[0].unsqueeze(0), boxes[2].unsqueeze(0)))
print('- clip_boxes_to_image (size = (100, 100)) :\n', clip_boxes_to_image(boxes, (100,100)))
print('- box_area :',box_area(boxes))
print('- remove_small_boxes (min_size = (60)) :', remove_small_boxes(boxes, 60))

cv2.imshow('test', black_canvas)
cv2.waitKey(0)
cv2.destroyAllWindows()

 

 

코드 결과