물체 탐지 :: 컴퓨터 비전 - mindscale
Skip to content

물체 탐지

실습 데이터 불러오기

이미지

from PIL import Image
import numpy as np

img = Image.open('bus_stop.jpg')

COCO 레이블

coco_labels = open('coco.txt').read().split()

SSD

텐서플로 허브에서 모형 불러오기

import tensorflow_hub as hub
detector = hub.load("https://tfhub.dev/tensorflow/ssd_mobilenet_v2/2")

모형에 이미지 입력

a = np.array(img) [np.newaxis]
output = detector(a)

물체 인식 결과 정리

boxes = output['detection_boxes'].numpy()[0]  # 경계 상자
labels = output['detection_classes'].numpy()[0] # 물체 종류
scores = output['detection_scores'].numpy()[0] # 점수
w, h = img.size  # 이미지 크기
results = []
for score, label, (y1, x1, y2, x2) in zip(scores, labels, boxes):
    results.append({
        'score': score,
        'label': coco_labels[int(label)],
        'box': {
            'xmin': int(w * x1),
            'ymin': int(h * y1),
            'xmax': int(w * x2),
            'ymax': int(h * y2)}
    })    

결과 시각화

시각화 함수

from PIL import ImageFont, ImageDraw, Image

def draw_result(img, results, threshold=.6, fontsize=30, color=(0, 255, 0)):
    img = img.copy()
    draw = ImageDraw.Draw(img)  # 이미지에 그림을 덧그리는 객체
    font = ImageFont.truetype("NanumGothic.ttf", fontsize)

    for result in results:  
        if result['score'] > threshold:
            box_pos = (                 # 사각형의
                result['box']['xmin'],  # 좌상단 좌표와
                result['box']['ymin'],
                result['box']['xmax'],  # 우하단 좌표를 지정
                result['box']['ymax'])

            draw.rectangle(box_pos, outline=color, width=3)

            text_pos = box_pos[:2]  # 좌상단 좌표를 텍스트의 위치로 지정
            draw.text(text_pos, result['label'], color, font=font)
    return img  # 이미지를 출력

시각화

draw_result(img, results)

DETR

transformers와 timm 패키지를 설치

pip install transformers timm

DETR 모델을 다운로드

from transformers import pipeline
detr = pipeline(model='facebook/detr-resnet-101')

모델에 이미지 입력

results = detr(img)

결과 시각화

draw_result(img, results)