ComfyUI-SelVA/selva_core/data/extraction/vgg_sound.py

import logging
import os
from pathlib import Path
from typing import Optional, Union

import pandas as pd
import torch
import torchaudio
from torch.utils.data.dataset import Dataset
from torchvision.transforms import v2
from torio.io import StreamingMediaDecoder

from selva_core.data.av_utils import normalize_video_chunk
from selva_core.utils.dist_utils import local_rank

log = logging.getLogger()

_CLIP_SIZE = 384
_CLIP_FPS = 8.0

_SYNC_SIZE = 224
_SYNC_FPS = 25.0


class VGGSound(Dataset):

    def __init__(
        self,
        root: Union[str, Path],
        *,
        tsv_path: Union[str, Path] = 'sets/vgg3-train.tsv',
        audio_required: bool = True,
        sample_rate: int = 16_000,
        duration_sec: float = 8.0,
        audio_samples: Optional[int] = None,
        normalize_audio: bool = False,
        clip_video_required: bool = True,
    ):
        self.root = Path(root)
        self.audio_required = audio_required
        if audio_required:
            self.normalize_audio = normalize_audio
            if audio_samples is None:
                self.audio_samples = int(sample_rate * duration_sec)
            else:
                self.audio_samples = audio_samples
                effective_duration = audio_samples / sample_rate
                # make sure the duration is close enough, within 15ms
                assert abs(effective_duration - duration_sec) < 0.015, \
                    f'audio_samples {audio_samples} does not match duration_sec {duration_sec}'
        self.clip_video_required = clip_video_required
        
        videos = sorted(os.listdir(self.root))
        videos = set([Path(v).stem for v in videos])  # remove extensions
        self.labels = {}
        self.videos = []
        missing_videos = []

        # read the tsv for subset information
        df_list = pd.read_csv(tsv_path, sep='\t', dtype={'id': str}).to_dict('records')
        for record in df_list:
            id = record['id']
            label = record['label']
            if id in videos:
                self.labels[id] = label
                self.videos.append(id)
            else:
                missing_videos.append(id)
        
        if local_rank == 0:
            log.info(f'{len(videos)} videos found in {root}')
            log.info(f'{len(self.videos)} videos found in {tsv_path}')
            log.info(f'{len(missing_videos)} videos missing in {root}')

        self.sample_rate = sample_rate
        self.duration_sec = duration_sec

        if audio_required:
            self.expected_audio_length = self.audio_samples
        self.sync_expected_length = int(_SYNC_FPS * self.duration_sec)
        if clip_video_required:
            self.clip_expected_length = int(_CLIP_FPS * self.duration_sec)

        self.sync_transform = v2.Compose([
            v2.Resize((_SYNC_SIZE, _SYNC_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
            # v2.CenterCrop(_SYNC_SIZE),
            v2.ToImage(),
            v2.ToDtype(torch.float32, scale=True),
            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])
        
        if clip_video_required:
            self.clip_transform = v2.Compose([
                v2.Resize((_CLIP_SIZE, _CLIP_SIZE), interpolation=v2.InterpolationMode.BICUBIC),
                v2.ToImage(),
                v2.ToDtype(torch.float32, scale=True),
            ])    
        if audio_required:
            self.resampler = {}

    def sample(self, idx: int) -> dict[str, torch.Tensor]:
        video_id = self.videos[idx]
        
        label = self.labels[video_id]
            
        reader = StreamingMediaDecoder(self.root / (video_id + '.mp4'))
        reader.add_basic_video_stream(
            frames_per_chunk=int(_SYNC_FPS * self.duration_sec),
            frame_rate=_SYNC_FPS,
            format='rgb24',
        )
        if self.audio_required:
            reader.add_basic_audio_stream(frames_per_chunk=2**30, )
        if self.clip_video_required:
            reader.add_basic_video_stream(
                frames_per_chunk=int(_CLIP_FPS * self.duration_sec),
                frame_rate=_CLIP_FPS,
                format='rgb24',
            )

        reader.fill_buffer()
        data_chunk = reader.pop_chunks()
        
        sync_chunk = data_chunk[0]
        if sync_chunk is None:
            raise RuntimeError(f'Sync video returned None {video_id}')
        sync_chunk = normalize_video_chunk(sync_chunk, self.sync_expected_length, 
                                           n_tolerance_frame=3, desc=video_id)
        sync_chunk = self.sync_transform(sync_chunk)
        
        if self.audio_required:
            audio_chunk = data_chunk[1]
        
        if self.clip_video_required:
            clip_chunk = data_chunk[2 if self.audio_required else 1]
            if clip_chunk is None:
                raise RuntimeError(f'CLIP video returned None {video_id}')
            clip_chunk = normalize_video_chunk(clip_chunk, self.clip_expected_length,
                                               n_tolerance_frame=1, desc=video_id)
            clip_chunk = self.clip_transform(clip_chunk)

        # process audio
        if self.audio_required:
            sample_rate = int(reader.get_out_stream_info(1).sample_rate)
            audio_chunk = audio_chunk.transpose(0, 1)
            audio_chunk = audio_chunk.mean(dim=0)  # mono
            if self.normalize_audio:
                abs_max = audio_chunk.abs().max()
                audio_chunk = audio_chunk * (0.95 / abs_max)
                if abs_max <= 1e-6:
                    raise RuntimeError(f'Audio is silent {video_id}')

            # resample
            if sample_rate == self.sample_rate:
                audio_chunk = audio_chunk
            else:
                if sample_rate not in self.resampler:
                    # https://pytorch.org/audio/stable/tutorials/audio_resampling_tutorial.html#kaiser-best
                    self.resampler[sample_rate] = torchaudio.transforms.Resample(
                        sample_rate,
                        self.sample_rate,
                        lowpass_filter_width=64,
                        rolloff=0.9475937167399596,
                        resampling_method='sinc_interp_kaiser',
                        beta=14.769656459379492,
                    )
                audio_chunk = self.resampler[sample_rate](audio_chunk)

            if audio_chunk.shape[0] < self.expected_audio_length:
                raise RuntimeError(f'Audio too short {video_id}')
            audio_chunk = audio_chunk[:self.expected_audio_length]

        data = {
            'id': video_id,
            'caption': label,
            'sync_video': sync_chunk,
        }
        
        if self.audio_required:
            data['audio'] = audio_chunk
        if self.clip_video_required:
            data['clip_video'] = clip_chunk

        return data

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        try:
            return self.sample(idx)
        except Exception as e:
            log.error(f'Error loading video {self.videos[idx]}: {e}')
            return None

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