[실전 예제/객체 추적/PyTorch] MOT 데이터셋으로 객체 추적 데이터셋 구성하기

객체 추적(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
import random
import argparse
from glob import glob
from collections import defaultdict
from tqdm import tqdm

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--seq_dir", required=True, help="예: ./MOT17/train/MOT17-02-FRCNN")
    p.add_argument("--out_dir", default="mot_crops", help="출력 폴더")
    p.add_argument("--crop_w", type=int, default=128)
    p.add_argument("--crop_h", type=int, default=256)
    p.add_argument("--pos_per_id", type=int, default=10)
    p.add_argument("--neg_per_id", type=int, default=10)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()

def main():
    args = parse_args()
    random.seed(args.seed)

    img_dir = os.path.join(args.seq_dir, "img1")
    gt_file = os.path.join(args.seq_dir, "gt", "gt.txt")

    if not os.path.isdir(img_dir):
        raise FileNotFoundError(f"[ERR] img1 not found: {img_dir}")
    if not os.path.isfile(gt_file):
        raise FileNotFoundError(f"[ERR] gt.txt not found: {gt_file}")

    out_dir = args.out_dir
    img_out_dir = os.path.join(out_dir, "images")
    pair_file = os.path.join(out_dir, "pairs.txt")
    os.makedirs(img_out_dir, exist_ok=True)

    # -------------------------
    # 1) GT 로드
    # gt format: frame,id,x,y,w,h,conf,class,visibility
    # -------------------------
    tracks = defaultdict(list)
    with open(gt_file, "r") as f:
        for line in f:
            vals = line.strip().split(",")
            if len(vals) < 6:
                continue
            frame = int(vals[0])
            tid = int(vals[1])
            x, y, w, h = map(float, vals[2:6])
            tracks[tid].append((frame, x, y, w, h))

    print(f"[INFO] seq_dir   : {args.seq_dir}")
    print(f"[INFO] img_dir   : {img_dir}")
    print(f"[INFO] gt_file   : {gt_file}")
    print(f"[INFO] IDs loaded: {len(tracks)}")

    # -------------------------
    # 2) Crop 생성
    # -------------------------
    id_to_imgs = defaultdict(list)

    for tid, boxes in tqdm(tracks.items(), desc="Cropping"):
        for frame, x, y, w, h in boxes:
            src = os.path.join(img_dir, f"{frame:06d}.jpg")
            if not os.path.isfile(src):
                # 일부 MOT는 png일 수도 있어서 fallback
                src_png = os.path.join(img_dir, f"{frame:06d}.png")
                src = src_png if os.path.isfile(src_png) else src

            img = cv2.imread(src)
            if img is None:
                continue

            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)

            # clamp
            H, W = img.shape[:2]
            x1 = max(0, min(W - 1, x1))
            y1 = max(0, min(H - 1, y1))
            x2 = max(0, min(W, x2))
            y2 = max(0, min(H, y2))
            if x2 <= x1 or y2 <= y1:
                continue

            crop = img[y1:y2, x1:x2]
            if crop.size == 0:
                continue

            crop = cv2.resize(crop, (args.crop_w, args.crop_h), interpolation=cv2.INTER_LINEAR)

            # 파일명에 seq 이름을 넣어 충돌 방지
            seq_name = os.path.basename(args.seq_dir)
            name = f"{seq_name}_{frame:06d}_{tid}.jpg"
            out_path = os.path.join(img_out_dir, name)
            cv2.imwrite(out_path, crop)
            id_to_imgs[tid].append(name)

    total_imgs = sum(len(v) for v in id_to_imgs.values())
    print(f"[INFO] crops saved: {total_imgs}")

    # -------------------------
    # 3) pairs.txt 생성
    # -------------------------
    ids = [tid for tid, imgs in id_to_imgs.items() if len(imgs) >= 2]
    if len(ids) < 2:
        raise RuntimeError("[ERR] pair를 만들기엔 ID가 너무 적습니다. (ID>=2, 각 ID 이미지>=2 필요)")

    pairs = []
    for tid in ids:
        imgs = id_to_imgs[tid]

        # positive
        for _ in range(args.pos_per_id):
            a, b = random.sample(imgs, 2)
            pairs.append(f"{a},{b},1")

        # negative
        neg_ids = [i for i in ids if i != tid]
        for _ in range(args.neg_per_id):
            nid = random.choice(neg_ids)
            a = random.choice(imgs)
            b = random.choice(id_to_imgs[nid])
            pairs.append(f"{a},{b},0")

    with open(pair_file, "w") as f:
        f.write("\n".join(pairs))

    print(f"[DONE] out_dir : {out_dir}")
    print(f"[DONE] images  : {img_out_dir}")
    print(f"[DONE] pairs   : {pair_file}  (lines={len(pairs)})")

if __name__ == "__main__":
    main()

 

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

사용법

python generate_mot_reid_dataset.py --seq_dir ./MOT17/train/MOT17-02-FRCNN --out_dir mot_crops

 

마무리

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

 

관련 내용

• [PyTorch] 맞춤형 데이터셋 만들기: torch.utils.data.Dataset() 사용 가이드
• [PyTorch] 효율적인 데이터 배치: torch.utils.data.DataLoader() 사용 가이드