[실전 예제/객체 탐지/PyTorch] DOTA 객체 검출 모델 구성과 학습

DOTA 기반 객체 검출(Object Detection)이란?

  객체 검출(Object Detection)은 이미지 내에서 객체의 위치와 클래스를 동시에 예측하는 컴퓨터 비전 태스크입니다.   DOTA(Dataset for Object Detection in Aerial Images)는 항공·위성 이미지 환경에서의 객체 검출을 목적으로 만들어진 대표적인 데이터셋으로, 일반적인 자연 이미지와 달리 작은 객체, 복잡한 배경, 고해상도 이미지라는 특징을 가집니다.

  특히 DOTA는 차량, 선박, 건물, 비행기 등 다양한 객체가 매우 작은 크기로 분포되어 있어 단순한 COCO 기반 학습 방식으로는 성능이 잘 나오지 않는 경우가 많습니다. 이번 글에서는 이전 글에서 구성한 DOTA 데이터셋을 기반으로, PyTorch를 이용해 객체 검출 모델을 구성하고 실제 학습까지 연결하는 과정을 정리합니다.

• 예: 항공 이미지에서 차량, 선박, 비행기 자동 검출
• 예: 위성 영상 기반 시설물 모니터링

 

목표

1. DOTA 데이터셋을 입력으로 받는 객체 검출 학습 파이프라인 구성
2. torchvision Faster R-CNN 기반 모델 구조 정의
3. 고해상도 이미지 환경에서 학습 시 주의할 점 정리

 

학습을 위한 데이터 구조

  이전 글에서 구성한 DOTA 데이터셋은 다음과 같은 구조를 가집니다. 이미지와 어노테이션이 COCO 형식(JSON)으로 변환되어 있어, torchvision Detection 모델과 바로 연결할 수 있습니다.

dota/
  train/
    images/
      xxx.jpg
    annotations.json
  val/
    images/
      xxx.jpg
    annotations.json

  DOTA는 원본 해상도가 매우 크기 때문에, 보통 이미지를 타일 단위로 분할한 뒤 학습에 사용합니다.

 

객체 검출 모델 구성

모델 개요

• Backbone: ResNet50 + FPN
• Detector: Faster R-CNN
• 입력: 타일링된 DOTA 이미지
• 출력: bounding box + class score

import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

def build_model(num_classes):
    model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(
        in_features,
        num_classes
    )
    return model

  DOTA 클래스 수에 맞게 마지막 box predictor만 교체하면, 나머지 구조는 그대로 재사용할 수 있습니다.

 

학습 코드

from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = build_model(num_classes=16)  # DOTA 클래스 수 + background
model.to(device)

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

for epoch in range(10):
    model.train()
    loss_sum = 0.0

    for images, targets in train_loader:
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        loss_sum += losses.item()

    print(f"[Epoch {epoch+1}] loss={loss_sum/len(train_loader):.4f}")

  DOTA 데이터는 작은 객체가 많아 초기 학습 시 loss 변동이 큽니다. 따라서 learning rate를 과하게 키우지 않고, 학습 로그를 충분히 확인하는 것이 중요합니다.

 

마무리

  이번 글에서는 DOTA 데이터셋을 기반으로 객체 검출 모델을 구성하고, PyTorch를 이용해 실제 학습까지 연결하는 과정을 살펴보았습니다. 항공 이미지 환경에서는 데이터 특성상 전처리와 타일링 전략이 모델 성능에 큰 영향을 미치므로, 데이터 구성 단계와 함께 고려하는 것이 중요합니다.

  다음 글에서는 학습된 모델을 이용해 실제 DOTA 이미지에 대한 추론 결과를 시각화하고, 작은 객체 검출 성능을 개선하기 위한 실전 팁을 정리해보겠습니다.

 

관련 내용

• [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: DOTA 데이터셋으로 PyTorch 데이터셋 만들기
• [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: COCO 데이터셋으로 PyTorch 데이터셋 만들기
• [실전 예제/PyTorch/Detection] YOLO 모델을 이용한 이미지 객체 탐지