반응형
DOTA 기반 객체 검출(Object Detection)이란?
객체 검출(Object Detection)은 이미지 내에서 객체의 위치와 클래스를 동시에 예측하는 컴퓨터 비전 태스크입니다. DOTA(Dataset for Object Detection in Aerial Images)는 항공·위성 이미지 환경에서의 객체 검출을 목적으로 만들어진 대표적인 데이터셋으로, 일반적인 자연 이미지와 달리 작은 객체, 복잡한 배경, 고해상도 이미지라는 특징을 가집니다.
특히 DOTA는 차량, 선박, 건물, 비행기 등 다양한 객체가 매우 작은 크기로 분포되어 있어 단순한 COCO 기반 학습 방식으로는 성능이 잘 나오지 않는 경우가 많습니다. 이번 글에서는 이전 글에서 구성한 DOTA 데이터셋을 기반으로, PyTorch를 이용해 객체 검출 모델을 구성하고 실제 학습까지 연결하는 과정을 정리합니다.
- 예: 항공 이미지에서 차량, 선박, 비행기 자동 검출
- 예: 위성 영상 기반 시설물 모니터링
목표
- DOTA 데이터셋을 입력으로 받는 객체 검출 학습 파이프라인 구성
- torchvision Faster R-CNN 기반 모델 구조 정의
- 고해상도 이미지 환경에서 학습 시 주의할 점 정리
학습을 위한 데이터 구조
이전 글에서 구성한 DOTA 데이터셋은 다음과 같은 구조를 가집니다. 이미지와 어노테이션이 COCO 형식(JSON)으로 변환되어 있어, torchvision Detection 모델과 바로 연결할 수 있습니다. DOTA 원본은 회전 박스(Oriented BBox)지만, 본 글에서는 COCO 형식의 axis-aligned bbox로 변환된 버전을 사용하므로, 회전 정보가 필요한 경우 별도의 OBB 모델이 필요합니다.
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
# -------------------------
# Model
# -------------------------
def build_model(num_classes: int):
try:
model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
except TypeError:
model = fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return modelDOTA 클래스 수에 맞게 마지막 box predictor만 교체하면, 나머지 구조는 그대로 재사용할 수 있습니다.
학습 코드
실행 예시
python train_fasterrcnn_dota_cocojson.py --root ./dota --epochs 10 --batch 2 --num_workers 4# train_fasterrcnn_dota_cocojson.py
import os
import json
import argparse
from typing import Any, Dict, List, Tuple
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision.transforms import functional as F
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
# 실행 예시
# python train_fasterrcnn_dota_cocojson.py --root ./dota --epochs 10 --batch 2 --num_workers 4
# -------------------------
# Utils
# -------------------------
def collate_fn(batch):
return tuple(zip(*batch))
def set_seed(seed: int):
import random
import numpy as np
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# -------------------------
# COCO-style JSON helpers
# -------------------------
def load_coco_json(path: str) -> dict:
if not os.path.isfile(path):
raise FileNotFoundError(f"[ERR] annotations not found: {path}")
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def build_category_mapping_from_cocojson(coco: dict) -> Tuple[Dict[int, int], int]:
"""
category_id(원본) -> contiguous label(1..K)
background=0
"""
cats = coco.get("categories", [])
if len(cats) == 0:
raise ValueError("[ERR] 'categories' is empty in annotations.json")
cat_ids = sorted(int(c["id"]) for c in cats)
cat_map = {cid: i + 1 for i, cid in enumerate(cat_ids)} # 1..K
num_classes = len(cat_ids) + 1
return cat_map, num_classes
# -------------------------
# Dataset: DOTA (COCO JSON converted)
# -------------------------
class DotaCocoJsonDetectionDataset(Dataset):
"""
DOTA를 COCO JSON 형식으로 변환해둔 annotations.json을 읽어서
torchvision Faster R-CNN 학습 입력 (image, target)을 반환.
Returns:
image: FloatTensor [3,H,W] in [0,1]
target:
boxes FloatTensor [N,4] xyxy
labels Int64Tensor [N] contiguous (1..K)
image_id Int64Tensor [1]
area FloatTensor [N] (optional but recommended)
iscrowd Int64Tensor [N]
"""
def __init__(self, images_dir: str, ann_file: str, cat_id_to_contiguous: Dict[int, int]):
self.images_dir = images_dir
self.ann_file = ann_file
self.cat_map = cat_id_to_contiguous
if not os.path.isdir(self.images_dir):
raise FileNotFoundError(f"[ERR] images_dir not found: {self.images_dir}")
coco = load_coco_json(self.ann_file)
self.id_to_img = {img["id"]: img for img in coco.get("images", [])}
self.img_ids = sorted(self.id_to_img.keys())
self.anns_by_img: Dict[int, List[dict]] = {}
for ann in coco.get("annotations", []):
img_id = ann["image_id"]
self.anns_by_img.setdefault(img_id, []).append(ann)
def __len__(self) -> int:
return len(self.img_ids)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, Any]]:
img_id = self.img_ids[idx]
info = self.id_to_img[img_id]
file_name = info["file_name"]
img_path = os.path.join(self.images_dir, file_name)
img = Image.open(img_path).convert("RGB")
image = F.to_tensor(img) # float32, [0,1]
H, W = image.shape[1], image.shape[2]
anns = self.anns_by_img.get(img_id, [])
boxes: List[List[float]] = []
labels: List[int] = []
areas: List[float] = []
iscrowd: List[int] = []
for ann in anns:
x, y, w, h = ann["bbox"] # COCO: [x,y,w,h]
if w <= 1 or h <= 1:
continue
x1, y1 = float(x), float(y)
x2, y2 = float(x + w), float(y + h)
# clamp (타일링/변환 과정에서 경계 조금 벗어나는 케이스 방어)
x1 = max(0.0, min(x1, W - 1.0))
y1 = max(0.0, min(y1, H - 1.0))
x2 = max(0.0, min(x2, float(W)))
y2 = max(0.0, min(y2, float(H)))
if x2 <= x1 or y2 <= y1:
continue
cat_id = int(ann["category_id"])
if cat_id not in self.cat_map:
continue
boxes.append([x1, y1, x2, y2])
labels.append(self.cat_map[cat_id])
# area는 COCO JSON에 있으면 사용, 없으면 bbox로 대체
areas.append(float(ann.get("area", (x2 - x1) * (y2 - y1))))
iscrowd.append(int(ann.get("iscrowd", 0)))
if len(boxes) == 0:
boxes_t = torch.zeros((0, 4), dtype=torch.float32)
labels_t = torch.zeros((0,), dtype=torch.int64)
areas_t = torch.zeros((0,), dtype=torch.float32)
iscrowd_t = torch.zeros((0,), dtype=torch.int64)
else:
boxes_t = torch.tensor(boxes, dtype=torch.float32)
labels_t = torch.tensor(labels, dtype=torch.int64)
areas_t = torch.tensor(areas, dtype=torch.float32)
iscrowd_t = torch.tensor(iscrowd, dtype=torch.int64)
target = {
"boxes": boxes_t,
"labels": labels_t,
"image_id": torch.tensor([img_id], dtype=torch.int64),
"area": areas_t,
"iscrowd": iscrowd_t,
}
return image, target
# -------------------------
# Model
# -------------------------
def build_model(num_classes: int):
try:
model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
except TypeError:
model = fasterrcnn_resnet50_fpn(pretrained=True)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
# -------------------------
# Train / quick val loss
# -------------------------
def train_one_epoch(model, loader, optimizer, device, epoch: int, print_every: int = 50):
model.train()
loss_sum = 0.0
for it, (images, targets) in enumerate(loader, start=1):
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(v for v in loss_dict.values())
optimizer.zero_grad(set_to_none=True)
losses.backward()
optimizer.step()
loss_sum += float(losses.item())
if it % print_every == 0:
avg = loss_sum / it
detail = {k: float(v.item()) for k, v in loss_dict.items()}
print(f"[Epoch {epoch}] iter={it}/{len(loader)} loss={avg:.4f} detail={detail}")
return loss_sum / max(1, len(loader))
@torch.no_grad()
def compute_loss_on_loader(model, loader, device):
# torchvision detection은 eval()에서 loss가 안 나오는 경우가 있어 train()+no_grad로 간단 계산
model.train()
loss_sum = 0.0
n = 0
for images, targets in 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)
loss_sum += float(sum(v for v in loss_dict.values()).item())
n += 1
return loss_sum / max(1, n)
def save_ckpt(model, optimizer, epoch: int, out_path: str):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
torch.save({"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict()}, out_path)
# -------------------------
# Main
# -------------------------
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--root", default="dota", help="dota root dir: train/val with images + annotations.json")
p.add_argument("--epochs", type=int, default=10)
p.add_argument("--batch", type=int, default=2)
p.add_argument("--lr", type=float, default=0.005)
p.add_argument("--momentum", type=float, default=0.9)
p.add_argument("--weight_decay", type=float, default=0.0005)
p.add_argument("--num_workers", type=int, default=4)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--device", default="cuda")
p.add_argument("--out_dir", default="runs/fasterrcnn_dota")
p.add_argument("--print_every", type=int, default=50)
return p.parse_args()
def main():
args = parse_args()
set_seed(args.seed)
device = torch.device(args.device if torch.cuda.is_available() and args.device.startswith("cuda") else "cpu")
train_images = os.path.join(args.root, "train", "images")
train_ann = os.path.join(args.root, "train", "annotations.json")
val_images = os.path.join(args.root, "val", "images")
val_ann = os.path.join(args.root, "val", "annotations.json")
train_coco = load_coco_json(train_ann)
cat_map, num_classes = build_category_mapping_from_cocojson(train_coco)
print(f"[INFO] root : {args.root}")
print(f"[INFO] num_classes : {num_classes} (incl background)")
print(f"[INFO] train_images: {train_images}")
print(f"[INFO] train_ann : {train_ann}")
print(f"[INFO] val_images : {val_images}")
print(f"[INFO] val_ann : {val_ann}")
train_ds = DotaCocoJsonDetectionDataset(train_images, train_ann, cat_map)
val_ds = DotaCocoJsonDetectionDataset(val_images, val_ann, cat_map)
train_loader = DataLoader(
train_ds,
batch_size=args.batch,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=collate_fn,
)
val_loader = DataLoader(
val_ds,
batch_size=args.batch,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=collate_fn,
)
model = build_model(num_classes)
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
print(f"[INFO] device={device} train={len(train_ds)} val={len(val_ds)}")
for epoch in range(1, args.epochs + 1):
train_loss = train_one_epoch(model, train_loader, optimizer, device, epoch, print_every=args.print_every)
val_loss = compute_loss_on_loader(model, val_loader, device)
lr_scheduler.step()
lr_now = optimizer.param_groups[0]["lr"]
print(f"[EPOCH {epoch}] train_loss={train_loss:.4f} val_loss={val_loss:.4f} lr={lr_now:.6f}")
ckpt_path = os.path.join(args.out_dir, f"epoch_{epoch:03d}.pth")
save_ckpt(model, optimizer, epoch, ckpt_path)
print(f"[DONE] checkpoints saved to: {args.out_dir}")
if __name__ == "__main__":
main()DOTA 데이터는 작은 객체가 많아 초기 학습 시 loss 변동이 큽니다. 따라서 learning rate를 과하게 키우지 않고, 학습 로그를 충분히 확인하는 것이 중요합니다.
마무리
이번 글에서는 DOTA 데이터셋을 기반으로 객체 검출 모델을 구성하고, PyTorch를 이용해 실제 학습까지 연결하는 과정을 살펴보았습니다. 항공 이미지 환경에서는 데이터 특성상 전처리와 타일링 전략이 모델 성능에 큰 영향을 미치므로, 데이터 구성 단계와 함께 고려하는 것이 중요합니다.
다음 글에서는 학습된 모델을 이용해 실제 DOTA 이미지에 대한 추론 결과를 시각화하고, 작은 객체 검출 성능을 개선하기 위한 실전 팁을 정리해보겠습니다.
관련 내용
- [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: DOTA 데이터셋으로 PyTorch 데이터셋 만들기
- [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: COCO 데이터셋으로 PyTorch 데이터셋 만들기
- [실전 예제/PyTorch/Detection] YOLO 모델을 이용한 이미지 객체 탐지
반응형
'실전 예제, 프로젝트' 카테고리의 다른 글
| [실전 예제/인스턴스 분할/PyTorch] Mask R-CNN 모델 구성과 COCO 학습 (0) | 2026.01.24 |
|---|---|
| [실전 예제/객체 탐지/PyTorch] Faster R-CNN 모델 구성과 COCO 학습 (0) | 2026.01.22 |
| [실전 예제/이미지 분류/PyTorch] ResNet 기반 이미지 분류 모델 구성과 학습 (0) | 2026.01.14 |
| [실전 예제/변화 탐지/PyTorch] Siamese 기반 변화 탐지 모델 구성과 학습 (0) | 2026.01.14 |
| [실전 예제/객체 추적/PyTorch] Re-ID 기반 객체 추적 모델 구성과 학습 (0) | 2026.01.14 |
