🐍 Python

허깅 페이스(Hugging Face)로 오브젝트 디텍션: 소화전 탐지 모델 만들기 A to Z

itstory(Booho) 2025. 2. 11. 10:35
728x90

1. 오브젝트 디텍션(Object Detection)이란?

오브젝트 디텍션은 이미지(또는 영상) 안에서 특정 물체가 어디에 있는지(바운딩 박스)와 무엇인지(클래스)를 찾아내는 기술입니다. 예를 들어, 사진 속 사람, 개, 자동차 등을 인식하고 위치를 표시해주는 기능이 대표적입니다.
최근에는 Faster R-CNN, YOLO, SSD 등 유명 모델에서 나아가, DETR, DINO 등 Transformer 기반 모델이 속속 등장하여 높은 정확도와 다양한 응용 가능성을 보여주고 있습니다.

2. 허깅 페이스(Hugging Face)에서의 오브젝트 디텍션

허깅 페이스는 원래 NLP(자연어 처리) 분야로 유명했지만, 지금은 컴퓨터 비전 모델도 활발하게 지원하고 있습니다.

  • DETR 시리즈
  • YOLOS, DINO
  • OWL-ViT (오픈 도메인 오브젝트 디텍션)
    등 다양한 모델이 허깅 페이스 모델 허브에 올라와 있어, 누구나 손쉽게 다운로드하고 테스트할 수 있습니다.

모델 허브에서 확인하기

허깅 페이스 모델 허브(https://huggingface.co/models)에서 object-detection으로 검색하면, 원하는 모델을 바로 찾아볼 수 있습니다.


3. 예시: DETR 모델 불러와서 추론하기

먼저, 허깅 페이스에 공개된 DETR(facebook/detr-resnet-50) 모델을 사용해 간단히 오브젝트 디텍션을 시도해봅시다.

pip install transformers torch torchvision
pip install matplotlib requests

(1) 기본 코드

import torch
import requests
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from transformers import DetrFeatureExtractor, DetrForObjectDetection

# 1) 이미지 불러오기
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tests/COCO_test.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# 2) 모델 및 특성 추출기 로드
model_name = "facebook/detr-resnet-50"
feature_extractor = DetrFeatureExtractor.from_pretrained(model_name)
model = DetrForObjectDetection.from_pretrained(model_name)

# 3) 전처리
inputs = feature_extractor(images=image, return_tensors="pt")

# 4) 추론(Inference)
outputs = model(**inputs)
results = feature_extractor.post_process(outputs, torch.Tensor([image.size[::-1]]))

# 5) 시각화
fig, ax = plt.subplots(1, figsize=(10, 6))
ax.imshow(image)

threshold = 0.7  # 신뢰도 기준
for result in results:
    scores = result["scores"].detach().cpu().numpy()
    labels = result["labels"].detach().cpu().numpy()
    boxes = result["boxes"].detach().cpu().numpy()  # [xmin, ymin, xmax, ymax]

    for score, label, box in zip(scores, labels, boxes):
        if score < threshold:
            continue

        xmin, ymin, xmax, ymax = box
        width, height = xmax - xmin, ymax - ymin
        rect = patches.Rectangle((xmin, ymin), width, height, 
                                 linewidth=2, edgecolor='r', facecolor='none')
        ax.add_patch(rect)
        ax.text(xmin, ymin, f"{label} {score:.2f}", 
                fontsize=12, color='white', 
                bbox=dict(facecolor='red', alpha=0.5))

plt.show()
  • DETR: 이미지 전체를 Transformer로 인코딩하여 오브젝트 위치와 클래스를 예측
  • FeatureExtractor: 이미지를 모델이 이해할 수 있는 형태로 전처리

코드를 실행하면, 사진에 검출된 물체들의 바운딩 박스가 그려집니다.


4. 소화전만 탐지하는 특화 모델 만들기

COCO 같은 일반 데이터셋에는 소화전(fire hydrant) 클래스도 포함되어 있습니다. 하지만 더 높은 정확도, 혹은 소화전에만 집중하여 불필요한 검출을 없애고 싶다면 직접 파인튜닝(fine-tuning) 해볼 수 있습니다.

4.1 데이터 수집 및 레이블링

  1. 소화전이 있는 이미지를 다양한 상황(거리 사진, 건물 앞 사진 등)에서 수집
  2. 각 이미지에서 소화전 위치를 바운딩 박스로 표시
  3. 레이블링 툴(CVAT, labelImg 등)을 사용해 COCO 포맷(혹은 Pascal VOC)으로 내보내기

COCO 포맷 예시(annotations.json):

{
  "images": [
    {"id": 1, "file_name": "image1.jpg", "height": 720, "width": 1280},
    ...
  ],
  "categories": [
    {"id": 1, "name": "fire hydrant"}
  ],
  "annotations": [
    {
      "id": 1,
      "image_id": 1,
      "category_id": 1,
      "bbox": [x, y, w, h],
      "area": 12345,
      "iscrowd": 0
    },
    ...
  ]
}

4.2 파인튜닝(재학습) 개요

  1. 이미 학습된 DETR 모델을 로드
  2. 내 데이터셋(소화전만 레이블링된 데이터)으로 재학습
  3. 모델이 소화전 클래스를 더욱 정확하게 인식하도록 만듦

파인튜닝에 필요한 주요 라이브러리:

pip install transformers datasets pycocotools

이후 transformers.Trainer나 직접 학습 루프를 작성해 진행할 수 있습니다. GPU 자원이 충분하다면, 구글 Colab 등에서 시도해보세요.


5. 추론 및 활용

파인튜닝이 완료된 모델을 이용해 새로운 이미지에서 소화전만 골라낼 수 있습니다. 가령, 다음처럼 테스트 이미지를 넣고 소화전 클래스(label=1)가 0.8 이상의 점수를 받는 경우만 따로 출력하도록 만들 수 있죠.

test_image = Image.open("some_street_view.jpg")
inputs = feature_extractor(images=test_image, return_tensors="pt")
outputs = model(**inputs)

results = feature_extractor.post_process(outputs, torch.Tensor([test_image.size[::-1]]))
for result in results:
    scores = result["scores"].detach().numpy()
    labels = result["labels"].detach().numpy()
    boxes = result["boxes"].detach().numpy()

    for score, label, box in zip(scores, labels, boxes):
        if score > 0.8 and label == 1:  # 1: 'fire hydrant'
            print("Fire hydrant found with score:", score)
            print("Bounding box:", box)

6. 모델 배포와 확장

  1. 허깅 페이스 모델 허브 업로드
    • transformers-cli login 후 trainer.push_to_hub() 활용
    • 협업자나 다른 환경에서 모델을 공유·다운로드 가능
  2. 웹 서비스로 배포
    • FastAPI, Flask 등을 통해 REST API 구축
    • 이미지 업로드 → 소화전 탐지 결과 반환 구조
  3. Edge/모바일 디바이스로 확장
    • 모델 경량화(ONNX, TensorRT) 후 임베디드 디바이스에 탑재
    • 현장 CCTV, 도시 안전 시스템 등 다양한 분야에 응용 가능

맺음말

  • 오브젝트 디텍션은 단순 이미지 분류를 넘어, 사진 속 여러 물체를 동시에 찾아내는 강력한 기술입니다.
  • 허깅 페이스에서는 DETR, DINO 등 모델이 이미 잘 정리되어 있어, 간단히 테스트하거나 파인튜닝하기가 수월합니다.
  • 더 나아가, 특정 클래스(예: 소화전)만 집중적으로 학습해 “내가 필요한 물체”를 더욱 정확하게 검출할 수 있습니다.
  • 완성된 모델은 다양한 형태(웹 서버, Edge 디바이스 등)로 배포할 수 있어, 실제 서비스에도 바로 활용 가능합니다.

이 글이 오브젝트 디텍션을 시작하고, 허깅 페이스를 통한 모델 파인튜닝 과정을 이해하는 데 도움이 되길 바랍니다. 궁금한 점은 언제든 댓글로 남겨주세요!


참고 자료

 

ㄴㄴ즐거운 딥러닝 & 컴퓨터 비전 탐험 되시길 바랍니다!

728x90