WLASL 수화 인식 실습 - VideoDataset, R3D, 학습과 추론

동영상에서 프레임을 뽑아 3D CNN으로 분류하기

동영상 데이터의 시공간 특성과 분석 기법(3D CNN, 2D CNN 변형, Attention), WLASL 데이터셋 소개는 앞선 글에서 다뤘다. 이번 글에서는 WLASL 어노테이션을 VideoDataset으로 읽고, Albumentations로 전처리한 뒤 R3D-18(3D CNN) 을 학습·검증·추론하는 흐름을 정리한다. 앞선 글에서 클래스 필터링까지 마친 annotations, class_list가 있다고 가정한다.

---

1. VideoDataset: 프레임 샘플링과 전처리

동영상 파일에서 고정 개수(예: 12장) 의 프레임을 골라 읽고, 정규화·텐서 변환까지 적용하는 Dataset이다.

• cv2.VideoCapture로 열고, num_frames 중 랜덤으로 num_to_select개 인덱스 선택.
• 프레임이 부족하면 마지막 프레임을 반복해 채움.
• 각 프레임을 중심 기준 정사각형 크롭(_crop_to_square) 후, Albumentations transform 적용.
• 반환: (frames, target) — frames는 (T, C, H, W) 형태 텐서, target은 클래스 인덱스.

import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset
from PIL import Image

class VideoDataset(Dataset):
    def __init__(self, video_root, annotations, class_list, transform=None):
        self.video_root = video_root
        self.annotations = annotations
        self.transform = transform
        self.class_list = class_list
        self.num_classes = len(class_list)

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

    def __getitem__(self, idx):
        annot = self.annotations[idx]
        video_path = os.path.join(self.video_root, annot['filename'])
        cap = cv2.VideoCapture(video_path)
        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        num_to_select = 12

        if num_frames < num_to_select:
            frame_idxs = list(range(num_frames)) + [num_frames - 1] * (num_to_select - num_frames)
        else:
            frame_idxs = np.sort(np.random.choice(num_frames, size=num_to_select, replace=False))

        frames = []
        for frame_idx in frame_idxs:
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = self._crop_to_square(frame)
            else:
                frame = np.array(Image.new('RGB', (224, 224), (0, 0, 0)))
            if self.transform:
                frame = self.transform(image=frame)['image']
            frames.append(frame)
        cap.release()

        frames = torch.stack(frames)
        target = torch.tensor(self.class_list.index(annot['label']), dtype=torch.long)
        return frames, target

    def _crop_to_square(self, image):
        h, w = image.shape[:2]
        s = min(h, w)
        cx, cy = w // 2, h // 2
        x = cx - s // 2
        y = cy - s // 2
        return image[y:y+s, x:x+s]

---

2. 학습/검증 분할

클래스별로 검증 집합에 최소 한 샘플이 들어가도록 하려면, 클래스당 하나씩 먼저 검증에 넣고 나머지를 학습에 둘 수 있다.

import random
random.shuffle(annotations)
train_annot, val_annot = [], []
val_cls_check = {}
for annot in annotations:
    if annot['label'] in val_cls_check:
        train_annot.append(annot)
    else:
        val_annot.append(annot)
        val_cls_check[annot['label']] = 1

---

3. 전처리: Albumentations

학습 시에는 리사이즈·패딩·정규화에 더해 ShiftScaleRotate, ColorJitter, RandomBrightnessContrast 등 증강을 넣고, 검증 시에는 리사이즈·패딩·정규화만 적용한다. ImageNet 평균·표준편차로 정규화한 뒤 ToTensorV2로 텐서로 만든다.

import albumentations as A
from albumentations.pytorch import ToTensorV2

hyper_params = {
    'num_epochs': 8,
    'lr': 0.0001,
    'image_size': 224,
    'train_batch_size': 12,
    'val_batch_size': 8,
}

train_transform = A.Compose([
    A.ShiftScaleRotate(rotate_limit=20, shift_limit=0.1, scale_limit=0.05, p=0.5, border_mode=0),
    A.ColorJitter(p=0.5),
    A.RandomBrightnessContrast(p=0.5),
    A.LongestMaxSize(max_size=hyper_params['image_size']),
    A.PadIfNeeded(min_height=hyper_params['image_size'], min_width=hyper_params['image_size'], border_mode=0),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), p=1.0),
    ToTensorV2()
])

val_transform = A.Compose([
    A.LongestMaxSize(max_size=hyper_params['image_size']),
    A.PadIfNeeded(min_height=hyper_params['image_size'], min_width=hyper_params['image_size'], border_mode=0),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), p=1.0),
    ToTensorV2()
])

train_dataset = VideoDataset(video_root=video_dir, annotations=train_annot, class_list=class_list, transform=train_transform)
val_dataset = VideoDataset(video_root=video_dir, annotations=val_annot, class_list=class_list, transform=val_transform)
train_dataloader = torch.utils.data.DataLoader(train_dataset, num_workers=2, batch_size=hyper_params['train_batch_size'], shuffle=True)
val_dataloader = torch.utils.data.DataLoader(val_dataset, num_workers=2, batch_size=hyper_params['val_batch_size'])

---

4. 3D CNN 모델: R3D-18

PyTorch torchvision.models.video.r3d_18은 3D ResNet으로, 입력 shape가 (N, C, T, H, W)(배치, 채널, 시간, 높이, 너비)이다. VideoDataset이 반환하는 것은 (N, T, C, H, W) 이므로, 학습·검증 시 permute(0, 2, 1, 3, 4) 로 바꿔 넣는다.

from torchvision import models
import torch.nn as nn

model = models.video.r3d_18(weights=models.video.R3D_18_Weights.DEFAULT)
model.fc = nn.Linear(512, len(class_list))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

---

5. 학습 루프

• 배치: (images, targets). images를 permute(0, 2, 1, 3, 4) 해서 (N, C, T, H, W)로 만든 뒤 모델에 넣는다.
• CrossEntropyLoss로 loss를 구하고, optimizer.step()으로 갱신.
• 검증에서는 같은 permute 후 예측하고, 정확도를 계산한 뒤, best 정확도일 때만 best_model.pth로 저장한다.

import numpy as np
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=hyper_params['lr'])
model_save_dir = './WLASL_train_results'
os.makedirs(model_save_dir, exist_ok=True)
best_acc = 0.0

for epoch in range(hyper_params['num_epochs']):
    model.train()
    for images, targets in train_dataloader:
        images = images.to(device).permute(0, 2, 1, 3, 4)
        targets = targets.to(device)
        outputs = model(images)
        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    model.eval()
    y_true, y_pred = [], []
    with torch.no_grad():
        for images, targets in val_dataloader:
            images = images.to(device).permute(0, 2, 1, 3, 4)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            y_true.extend(targets.cpu().numpy().tolist())
            y_pred.extend(preds.cpu().numpy().tolist())
    acc = (np.array(y_pred) == np.array(y_true)).mean()
    if acc > best_acc:
        torch.save(model.state_dict(), os.path.join(model_save_dir, 'best_model.pth'))
        best_acc = acc

---

6. 추론

새 동영상에서 12프레임을 같은 방식으로 샘플링하고, val_transform 적용 후 (1, C, T, H, W) 로 permute 해서 모델에 넣는다. F.softmax(output, dim=1)로 클래스별 확률을 구할 수 있다.

import torch.nn.functional as F

model.eval()
with torch.no_grad():
    frames = frames.unsqueeze(0).permute(0, 2, 1, 3, 4)
    output = model(frames.to(device))
    probs = F.softmax(output, dim=1)[0].cpu().numpy()
for cls, p in zip(class_list, probs):
    print(f"{cls}: {p*100:.2f}%")

---

7. 정리

주제	핵심 포인트
VideoDataset	동영상에서 고정 개수 프레임을 샘플링하고, 크롭·전처리 후 (T, C, H, W) 텐서로 반환한다.
R3D 입력	PyTorch 3D 비디오 모델은 (N, C, T, H, W)를 기대하므로, DataLoader 출력에 permute(0, 2, 1, 3, 4)를 적용한다.
학습/검증 분할	클래스당 한 샘플씩 검증에 두면 모든 클래스가 검증에 포함된다.
추론	동일한 프레임 수·전처리·permute를 적용한 뒤 softmax로 클래스별 확률을 얻는다.

---

마치며

• WLASL과 R3D-18을 이용하면 수화 단어 분류를 VideoDataset 구성부터 학습·추론까지 한 흐름으로 구현할 수 있다.
• 프레임 수, 샘플링 방식, 증강, 모델(R3D / I3D / 비디오 Transformer 등)을 바꿔 가며 실험하면 동영상 태스크에 맞는 설정을 찾기 쉽다.