[실전 예제/객체 추적/PyTorch] 객체 추적 튜토리얼: MOT 데이터셋으로 PyTorch 데이터셋 만들기

객체 추적(Object Tracking)이란?

  객체 추적(Object Tracking)은 영상 속에서 특정 객체를 시간에 따라 지속적으로 추적하는 컴퓨터 비전 태스크입니다. 프레임마다 객체의 위치를 예측하고, 각 객체에 대해 고유한 ID를 유지하는 것이 핵심이죠. 이번 시간에는 PyTorch를 이용하여 객체 추적 데이터셋을 만드는 방법에 대해 알아보도록 하겠습니다.

• 예: CCTV에서 사람을 추적하거나, 자율주행 차량에서 차량/보행자를 추적하는 경우

 

목표

1. 실제 MOT17 구조를 탐색해서 하위 시퀀스들을 하나로 합쳐서 Dataset 구성
2. 객체 Crop 이미지를 만들고, 이를 기반으로 Re-ID 임베딩 모델을 학습할 수 있도록 구성
3. PyTorch Dataset ↔ 추적 모델 학습 코드가 자연스럽게 연결되도록 통일 (2장에서 준비)

 

PyTorch용 MOT Object Tracking Dataset 만들기

MOT 라벨 포맷 (.txt)

• frame_id: 프레임 번호
• object_id: 객체의 고유 ID
• x, y, w, h: 바운딩 박스 정보
• conf: confidence (GT의 경우 보통 1.0)
• class_id: 객체 클래스 (사람: 1)
• visibility: 가려짐 정도 (0~1)

frame_id, object_id, x, y, w, h, conf, class_id, visibility

디렉토리 구조 예시

MOT17/
  train/
    MOT17-02-DPM/
      img1/
        000001.jpg
        ...
      gt/
        gt.txt
    MOT17-04-FRCNN/
      ...
  test/
    ...

 

PyTorch 코드 예제

import os
import cv2
from PIL import Image

import torch
from torch.utils.data import Dataset
import torchvision.transforms as T

import pandas as pd
from tqdm import tqdm


def generate_mot_crops(mot_root, output_dir, min_vis=0.0, size=(256, 128), pair_txt_name="pairs.txt"):
    os.makedirs(output_dir, exist_ok=True)
    pair_file = open(os.path.join(output_dir, pair_txt_name), 'w')
    pair_set = set()

    seq_dirs = [d for d in os.listdir(mot_root) if os.path.isdir(os.path.join(mot_root, d))]
    for seq in seq_dirs:
        seq_path = os.path.join(mot_root, seq)
        img_dir = os.path.join(seq_path, 'img1')
        gt_path = os.path.join(seq_path, 'gt', 'gt.txt')
        if not os.path.exists(gt_path):
            continue

        df = pd.read_csv(gt_path, header=None)
        df.columns = ['frame', 'id', 'x', 'y', 'w', 'h', 'conf', 'class', 'vis']
        df = df[df['conf'] == 1]
        df = df[df['vis'] >= min_vis]
        df = df[df['class'] == 1]  # person class

        id_to_imgs = {}

        for _, row in tqdm(df.iterrows(), total=len(df), desc=seq):
            frame = int(row['frame'])
            track_id = int(row['id'])
            img_path = os.path.join(img_dir, f"{frame:06d}.jpg")
            image = cv2.imread(img_path)
            H, W = image.shape[:2]
            # 좌표 보정
            x1 = max(0, int(row['x']))
            y1 = max(0, int(row['y']))
            x2 = min(W, x1 + int(row['w']))
            y2 = min(H, y1 + int(row['h']))
            if x2 <= x1 or y2 <= y1:
                continue  # 잘못된 박스는 스킵
            crop = image[y1:y2, x1:x2]
            if crop.size == 0:
                continue  # 혹시나 여전히 비면 스킵
            if image is None:
                continue
            crop = cv2.resize(crop, size)
            out_name = f"{seq}_{frame:06d}_{track_id}.jpg"
            out_path = os.path.join(output_dir, out_name)
            cv2.imwrite(out_path, crop)

            id_to_imgs.setdefault(track_id, []).append(out_name)

        # positive pairs
        for img_list in id_to_imgs.values():
            for i in range(len(img_list)):
                for j in range(i+1, len(img_list)):
                    pair = (img_list[i], img_list[j], 1)
                    pair_file.write(','.join(map(str, pair)) + '\n')
                    pair_set.add((img_list[i], img_list[j]))

        # negative pairs (한 시퀀스 내에서)
        ids = list(id_to_imgs.keys())
        for i in range(len(ids)):
            for j in range(i+1, len(ids)):
                img1 = id_to_imgs[ids[i]][0]
                img2 = id_to_imgs[ids[j]][0]
                if (img1, img2) not in pair_set:
                    pair_file.write(f"{img1},{img2},0\n")

    pair_file.close()

class MOTCropPairDataset(Dataset):
    def __init__(self, crop_dir, pair_file, transform=None):
        self.crop_dir = crop_dir
        self.transform = transform or T.Compose([
            T.ToTensor()
        ])
        self.pairs = []
        with open(pair_file, 'r') as f:
            for line in f:
                a, b, label = line.strip().split(',')
                self.pairs.append((a, b, int(label)))

    def __getitem__(self, idx):
        img1, img2, label = self.pairs[idx]
        path1 = os.path.join(self.crop_dir, img1)
        path2 = os.path.join(self.crop_dir, img2)
        img1 = Image.open(path1).convert('RGB')
        img2 = Image.open(path2).convert('RGB')
        return self.transform(img1), self.transform(img2), torch.tensor(label, dtype=torch.float32)

    def __len__(self):
        return len(self.pairs)

 

실제 MOT17 학습용 crop 데이터 자동 생성기

사용법

generate_mot_crops("MOT17/train", "mot_crops", min_vis=0.3)

 

마무리

  PyTorch를 이용하여 객체 추적 데이터셋을 어떻게 만드는지 살펴보았습니다. 다음 시간에는 모델 구성 및 학습 방법을 PyTorch로 작성하는 방법을 알아보도록 하겠습니다.

 

관련 내용

• 준비중