[실전 예제/객체 추적/PyTorch] Re-ID 기반 객체 추적 모델 구성과 학습

Re-ID 기반 객체 추적(Object Tracking)이란?

  객체 추적에서 가장 어려운 문제는 객체가 잠시 가려지거나 화면에서 사라졌다가 다시 등장했을 때도 동일한 ID를 유지하는 것입니다. 이를 위해 대부분의 현대적인 객체 추적 시스템은 Re-ID(Re-Identification) 모델을 사용합니다.

  Re-ID는 객체의 외형 정보를 임베딩 벡터로 변환하고, 시간적으로 떨어진 프레임 간에도 같은 객체인지 판단할 수 있도록 도와줍니다. 이번 글에서는 이전 글에서 생성한 MOT crop 데이터셋을 기반으로 Re-ID 임베딩 모델을 구성하고 학습하는 방법을 살펴봅니다.

• 예: 사람이 가려졌다가 다시 등장했을 때 같은 ID로 복원
• 예: 다른 사람과 교차한 이후에도 ID 스위치 방지

 

목표

1. MOT crop + pair 데이터셋을 이용한 Re-ID 학습 파이프라인 구성
2. Siamese 구조 기반 임베딩 모델 정의
3. 추후 Kalman Filter + Hungarian matching과 자연스럽게 결합 가능하도록 설계

 

Re-ID 학습을 위한 데이터 구조

이전 글에서 생성한 데이터는 다음과 같은 구조를 가집니다.

mot_crops/
  MOT17-02_000001_3.jpg
  MOT17-02_000002_3.jpg
  ...
  pairs.txt

pairs.txt는 두 이미지와 라벨로 구성된 pair 데이터입니다.

img_a.jpg,img_b.jpg,1   # 같은 객체
img_c.jpg,img_d.jpg,0   # 다른 객체

이 구조는 Contrastive Loss / Siamese Network 학습에 바로 사용할 수 있습니다.

 

Re-ID 임베딩 모델 구성

모델 개요

• Backbone: ResNet18
• 입력: 객체 crop 이미지 (256×128)
• 출력: L2 정규화된 임베딩 벡터

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet18


class ReIDEmbeddingNet(nn.Module):
    def __init__(self, embedding_dim=128):
        super().__init__()
        backbone = resnet18(pretrained=True)
        backbone.fc = nn.Identity()
        self.backbone = backbone

        self.head = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, embedding_dim)
        )

    def forward(self, x):
        feat = self.backbone(x)
        emb = self.head(feat)
        emb = F.normalize(emb, p=2, dim=1)
        return emb

  출력 임베딩은 L2 정규화를 거쳐 거리 기반 비교에 바로 사용할 수 있도록 합니다.

 

Contrastive Loss 정의

class ContrastiveLoss(nn.Module):
    def __init__(self, margin=1.0):
        super().__init__()
        self.margin = margin

    def forward(self, e1, e2, label):
        dist = torch.norm(e1 - e2, dim=1)
        pos = label * dist.pow(2)
        neg = (1 - label) * torch.clamp(self.margin - dist, min=0).pow(2)
        return (pos + neg).mean()

  같은 객체는 거리를 줄이고, 다른 객체는 margin 이상 벌어지도록 학습됩니다.

 

Re-ID 모델 학습 코드

from torch.utils.data import DataLoader

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

dataset = MOTCropPairDataset(
    crop_dir="mot_crops",
    pair_file="mot_crops/pairs.txt"
)

loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4)

model = ReIDEmbeddingNet(embedding_dim=128).to(device)
criterion = ContrastiveLoss(margin=1.0)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

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

    for img1, img2, label in loader:
        img1 = img1.to(device)
        img2 = img2.to(device)
        label = label.to(device)

        e1 = model(img1)
        e2 = model(img2)
        loss = criterion(e1, e2, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        loss_sum += loss.item()

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

 

 

마무리

  이번 글에서는 MOT 데이터셋에서 생성한 crop 데이터를 활용하여 Re-ID 임베딩 모델을 구성하고 학습하는 전체 과정을 살펴보았습니다. 이 모델은 이후 Kalman Filter, IOU, Hungarian Matching과 결합되어 실제 객체 추적 시스템의 핵심 구성 요소로 사용됩니다.

  다음 글에서는 학습된 Re-ID 임베딩을 이용해 프레임 간 객체를 매칭하고, 간단한 다중 객체 추적 파이프라인을 PyTorch로 구현해보겠습니다.

 

관련 내용

• [실전 예제/객체 추적/PyTorch] 객체 추적 튜토리얼: MOT 데이터셋으로 PyTorch 데이터셋 만들기
• [PyTorch] 맞춤형 데이터셋 만들기: torch.utils.data.Dataset() 사용 가이드
• [PyTorch] 효율적인 데이터 배치: torch.utils.data.DataLoader() 사용 가이드