data_utils/dataloader_torch.py

import sys
import os
sys.path.append(os.getcwd())
import os
from tqdm import tqdm
from data_utils.utils import *
import torch.utils.data as data
from data_utils.mesh_dataset import SmplxDataset
from transformers import Wav2Vec2Processor


class MultiVidData():
    def __init__(self, 
                data_root, 
                speakers, 
                split='train', 
                limbscaling=False, 
                normalization=False,
                norm_method='new',
                split_trans_zero=False,
                num_frames=25,
                num_pre_frames=25,
                num_generate_length=None,
                aud_feat_win_size=None,
                aud_feat_dim=64,
                feat_method='mel_spec',
                context_info=False,
                smplx=False,
                audio_sr=16000,
                convert_to_6d=False,
                expression=False,
                config=None
                ):
        self.data_root = data_root
        self.speakers = speakers
        self.split = split
        if split == 'pre':
            self.split = 'train'
        self.norm_method=norm_method
        self.normalization = normalization
        self.limbscaling = limbscaling
        self.convert_to_6d = convert_to_6d
        self.num_frames=num_frames
        self.num_pre_frames=num_pre_frames
        if num_generate_length is None:
            self.num_generate_length = num_frames
        else:
            self.num_generate_length = num_generate_length
        self.split_trans_zero=split_trans_zero

        dataset = SmplxDataset
        
        if self.split_trans_zero:
            self.trans_dataset_list = []
            self.zero_dataset_list = []
        else:
            self.all_dataset_list = []
        self.dataset={}
        self.complete_data=[]
        self.config=config
        load_mode=self.config.dataset_load_mode
        
        ######################load with pickle file
        if load_mode=='pickle':
            import pickle
            import subprocess
            
            # store_file_path='/tmp/store.pkl'
            # cp /is/cluster/scratch/hyi/ExpressiveBody/SMPLifyX4/scripts/store.pkl /tmp/store.pkl
            # subprocess.run(f'cp /is/cluster/scratch/hyi/ExpressiveBody/SMPLifyX4/scripts/store.pkl {store_file_path}',shell=True)
            
            # f = open(self.config.store_file_path, 'rb+')
            f = open(self.split+config.Data.pklname, 'rb+')
            self.dataset=pickle.load(f)
            f.close()
            for key in self.dataset:
                self.complete_data.append(self.dataset[key].complete_data)
        ######################load with pickle file
                
        ######################load with a csv file
        elif load_mode=='csv':

            # 这里从我的一个code文件夹导入的，后续再完善进来
            try:
                sys.path.append(self.config.config_root_path)
                from config import config_path
                from csv_parser import csv_parse
                
            except ImportError as e:
                print(f'err: {e}')
                raise ImportError('config root path error...')


            for speaker_name in self.speakers:
                # df_intervals=pd.read_csv(self.config.voca_csv_file_path)
                df_intervals=None
                df_intervals=df_intervals[df_intervals['speaker']==speaker_name]
                df_intervals = df_intervals[df_intervals['dataset'] == self.split]

                print(f'speaker {speaker_name} train interval length: {len(df_intervals)}')
                for iter_index, (_, interval) in tqdm(
                        (enumerate(df_intervals.iterrows())),desc=f'load {speaker_name}'
                ):
                    
                    (
                        interval_index,
                        interval_speaker,
                        interval_video_fn,
                        interval_id,
                        
                        start_time,
                        end_time,
                        duration_time,
                        start_time_10,
                        over_flow_flag,
                        short_dur_flag,
                        
                        big_video_dir,
                        small_video_dir_name,
                        speaker_video_path,
                        
                        voca_basename,
                        json_basename,
                        wav_basename,
                        voca_top_clip_path,
                        voca_json_clip_path,
                        voca_wav_clip_path,
                        
                        audio_output_fn,
                        image_output_path,
                        pifpaf_output_path,
                        mp_output_path,
                        op_output_path,
                        deca_output_path,
                        pixie_output_path,
                        cam_output_path, 
                        ours_output_path,
                        merge_output_path,
                        multi_output_path,
                        gt_output_path,
                        ours_images_path,
                        pkl_fil_path,
                    )=csv_parse(interval)
                    
                    if not os.path.exists(pkl_fil_path) or not os.path.exists(audio_output_fn):
                        continue

                    key=f'{interval_video_fn}/{small_video_dir_name}'
                    self.dataset[key] = dataset(
                        data_root=pkl_fil_path,
                        speaker=speaker_name,
                        audio_fn=audio_output_fn,
                        audio_sr=audio_sr,
                        fps=num_frames,
                        feat_method=feat_method,
                        audio_feat_dim=aud_feat_dim,
                        train=(self.split == 'train'),
                        load_all=True,
                        split_trans_zero=self.split_trans_zero,
                        limbscaling=self.limbscaling,
                        num_frames=self.num_frames,
                        num_pre_frames=self.num_pre_frames,
                        num_generate_length=self.num_generate_length,
                        audio_feat_win_size=aud_feat_win_size,
                        context_info=context_info,
                        convert_to_6d=convert_to_6d,
                        expression=expression,
                        config=self.config
                    )
                    self.complete_data.append(self.dataset[key].complete_data)
        ######################load with a csv file
                
        ######################origin load method
        elif load_mode=='json':
            
            # if self.split == 'train':
            #     import pickle
            #     f = open('store.pkl', 'rb+')
            #     self.dataset=pickle.load(f)
            #     f.close()
            #     for key in self.dataset:
            #         self.complete_data.append(self.dataset[key].complete_data)
            # else:https://pytorch-tutorial-assets.s3.amazonaws.com/VOiCES_devkit/source-16k/train/sp0307/Lab41-SRI-VOiCES-src-sp0307-ch127535-sg0042.wav
            # if config.Model.model_type == 'face':
            am = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-phoneme")
            am_sr = 16000
            # else:
            #     am, am_sr = None, None
            for speaker_name in self.speakers:
                speaker_root = os.path.join(self.data_root, speaker_name)

                videos=[v for v in os.listdir(speaker_root) ]
                print(videos)

                haode = huaide = 0

                for vid in tqdm(videos, desc="Processing training data of {}......".format(speaker_name)):
                    source_vid=vid
                    # vid_pth=os.path.join(speaker_root, source_vid, 'images/half', self.split)
                    vid_pth = os.path.join(speaker_root, source_vid, self.split)
                    if smplx == 'pose':
                        seqs = [s for s in os.listdir(vid_pth) if (s.startswith('clip'))]
                    else:
                        try:
                            seqs = [s for s in os.listdir(vid_pth)]
                        except:
                            continue

                    for s in seqs:
                        seq_root=os.path.join(vid_pth, s)
                        key = seq_root # correspond to clip******
                        audio_fname = os.path.join(speaker_root, source_vid, self.split, s, '%s.wav' % (s))
                        motion_fname = os.path.join(speaker_root, source_vid, self.split, s, '%s.pkl' % (s))
                        if not os.path.isfile(audio_fname) or not os.path.isfile(motion_fname):
                            huaide = huaide + 1
                            continue

                        self.dataset[key]=dataset(
                            data_root=seq_root,
                            speaker=speaker_name,
                            motion_fn=motion_fname,
                            audio_fn=audio_fname,
                            audio_sr=audio_sr,
                            fps=num_frames,
                            feat_method=feat_method,
                            audio_feat_dim=aud_feat_dim,
                            train=(self.split=='train'),
                            load_all=True,
                            split_trans_zero=self.split_trans_zero,
                            limbscaling=self.limbscaling,
                            num_frames=self.num_frames,
                            num_pre_frames=self.num_pre_frames,
                            num_generate_length=self.num_generate_length,
                            audio_feat_win_size=aud_feat_win_size,
                            context_info=context_info,
                            convert_to_6d=convert_to_6d,
                            expression=expression,
                            config=self.config,
                            am=am,
                            am_sr=am_sr,
                            whole_video=config.Data.whole_video
                        )
                        self.complete_data.append(self.dataset[key].complete_data)
                        haode = haode + 1
                print("huaide:{}, haode:{}".format(huaide, haode))
            import pickle

            f = open(self.split+config.Data.pklname, 'wb')
            pickle.dump(self.dataset, f)
            f.close()
        ######################origin load method

        self.complete_data=np.concatenate(self.complete_data, axis=0)

        # assert self.complete_data.shape[-1] == (12+21+21)*2
        self.normalize_stats = {}

        self.data_mean = None
        self.data_std = None
    
    def get_dataset(self):
        self.normalize_stats['mean'] = self.data_mean
        self.normalize_stats['std'] = self.data_std

        for key in list(self.dataset.keys()):
            if self.dataset[key].complete_data.shape[0] < self.num_generate_length:
                continue
            self.dataset[key].num_generate_length = self.num_generate_length
            self.dataset[key].get_dataset(self.normalization, self.normalize_stats, self.split)
            self.all_dataset_list.append(self.dataset[key].all_dataset)
        
        if self.split_trans_zero:
            self.trans_dataset = data.ConcatDataset(self.trans_dataset_list)
            self.zero_dataset = data.ConcatDataset(self.zero_dataset_list)
        else:
            self.all_dataset = data.ConcatDataset(self.all_dataset_list)