Pytorch实现姿态识别（一）——将视频数据处理为图像

ゞ浴缸里的玫瑰 2023-02-28 05:47 1阅读 0赞

### 文章目录 ###

*  一、mypath.py
 *  二、datasets.py

# 一、mypath.py #

class Path(object):
        @staticmethod #静态方法无需实例化被调用/也可以实例化后调用
        def db_dir(database):
            if database == 'ucf101':
                # folder that contains class labels
                root_dir = r'F:\DATASET\UCF-101\UCF-101'
    
                # Save preprocess data into output_dir
                output_dir = r'F:\DATASET\UCF-101\mydatasets'
    
                return root_dir, output_dir
            elif database == 'hmdb51':
                # folder that contains class labels
                root_dir = '/Path/to/hmdb-51'
    
                output_dir = '/path/to/VAR/hmdb51'
    
                return root_dir, output_dir
            else:
                print('Database {} not available.'.format(database))
                raise NotImplementedError
    
        @staticmethod
        def model_dir():
            return './model/c3d-pretrained.pth'

# 二、datasets.py #

以ucf-101数据集为例，将其视频数据转为图像数据，  
每个视频每隔4帧取到一张图像，一次读取16帧图像，

import os
    from sklearn.model_selection import train_test_split
    
    import torch
    import cv2
    import numpy as np
    from torch.utils.data import Dataset
    from mypath import Path
    
    class VideoDataset(Dataset):
        # 注意第一次要预处理数据的,开始时preprocess设置为True，训练时设置为False
        def __init__(self, dataset='ucf101', split='train', clip_len=16, preprocess=False):
            self.root_dir, self.output_dir = Path.db_dir(dataset)
            folder = os.path.join(self.output_dir, split)# 遍历train文件夹
            self.clip_len = clip_len# 一次读取多少帧图像当做序列
            self.split = split#
            self.resize_height = 128# 预处理视频高
            self.resize_width = 171# 预处理视频宽
            self.crop_size = 112# 预处理裁剪
    
            if not self.check_integrity():
                raise RuntimeError('Dataset not found or corrupted.' +
                                   ' You need to download it from official website.')
    
            if (not self.check_preprocess()) or preprocess:
                print('Preprocessing of {} dataset, this will take long, but it will be done only once.'.format(dataset))
                self.preprocess()
    
            self.fnames, labels = [], []
            for label in sorted(os.listdir(folder)):
                for fname in os.listdir(os.path.join(folder, label)):
                    self.fnames.append(os.path.join(folder, label, fname))
                    labels.append(label)
            # assert（断言）用于判断一个表达式，在表达式条件为 false 的时候触发异常。
            assert len(labels) == len(self.fnames)
            print('Number of {} videos: {:d}'.format(split, len(self.fnames)))
            """
            创建标签与索引之间的映射关系
            set返回的是{}，内建函数 sorted 方法返回的是一个新的 list
            返回的例子eg:{'ApplyEyeMakeup': 0, 'ApplyLipstick': 1,...}
            """
            self.label2index = {
        label: index for index, label in enumerate(sorted(set(labels)))}
            # 列表变ndarray，返回的是：[  0   0   0 ... 100 100 100]
            self.label_array = np.array([self.label2index[label] for label in labels], dtype=int)
    
            if dataset == "ucf101":
                if not os.path.exists('ucf_labels.txt'):
                    with open('ucf_labels.txt', 'w') as f:
                        for id, label in enumerate(sorted(self.label2index)):
                            f.writelines(str(id+1) + ' ' + label + '\n')
    
            elif dataset == 'hmdb51':
                if not os.path.exists('hmdb_labels.txt'):
                    with open('hmdb_labels.txt', 'w') as f:
                        for id, label in enumerate(sorted(self.label2index)):
                            f.writelines(str(id+1) + ' ' + label + '\n')
    
        def __len__(self):
            return len(self.fnames)
    
        #需要重写__getitem__方法
        def __getitem__(self, index):
            # 加载数据和预处理
            buffer = self.load_frames(self.fnames[index]) # 拿到每个文件夹,eg:v_ApplyEyeMakeup_g01_c01
            buffer = self.crop(buffer, self.clip_len, self.crop_size)# 不满足16帧的处理一下，得到数据X
            labels = np.array(self.label_array[index])# 标签
            if self.split == 'test':
                # 数据扩充
                buffer = self.randomflip(buffer)
            buffer = self.normalize(buffer)# 去均值
            buffer = self.to_tensor(buffer)# 网络的输入c×l×h×w（通道×帧数×高×宽）
            return torch.from_numpy(buffer), torch.from_numpy(labels)# ndarray转tensor
    
        def check_integrity(self):
            if not os.path.exists(self.root_dir):
                return False
            else:
                return True
    
        def check_preprocess(self):
            # TODO: Check image size in output_dir
            if not os.path.exists(self.output_dir):
                return False
            elif not os.path.exists(os.path.join(self.output_dir, 'train')):
                return False
    
            for ii, video_class in enumerate(os.listdir(os.path.join(self.output_dir, 'train'))):
                for video in os.listdir(os.path.join(self.output_dir, 'train', video_class)):
                    # 拿到每个类别的第一张图像：eg:F:\DATASET\UCF-101\xx\train\ApplyEyeMakeup\v_ApplyEyeMakeup_g01_c01\00000.jpg
                    video_name = os.path.join(os.path.join(self.output_dir, 'train', video_class, video),
                                        sorted(os.listdir(os.path.join(self.output_dir, 'train', video_class, video)))[0])
                    print(video_name)
                    image = cv2.imread(video_name)# 读出来的格式是HWC
                    if np.shape(image)[0] != 128 or np.shape(image)[1] != 171:
                        return False
                    else:
                        break
                # 10个类别就跳出
                if ii == 10:
                    break
    
            return True
    
        def preprocess(self):
            # 创建三个文件夹
            if not os.path.exists(self.output_dir):
                os.mkdir(self.output_dir)
                os.mkdir(os.path.join(self.output_dir, 'train'))
                os.mkdir(os.path.join(self.output_dir, 'val'))
                os.mkdir(os.path.join(self.output_dir, 'test'))
    
            for file in os.listdir(self.root_dir):# 遍历类别/标签
                file_path = os.path.join(self.root_dir, file)# 类别
                video_files = [name for name in os.listdir(file_path)]# 拿到所有视频
                # 视频分成3部分，sklearn只有能分训练集和测试集，我们把它训练集再分为训练集和验证集
                train_and_valid, test = train_test_split(video_files, test_size=0.2, random_state=42)
                train, val = train_test_split(train_and_valid, test_size=0.2, random_state=42)
                # 得到三个训练、测试、验证的文件路径
                train_dir = os.path.join(self.output_dir, 'train', file)
                val_dir = os.path.join(self.output_dir, 'val', file)
                test_dir = os.path.join(self.output_dir, 'test', file)
                # 不存在即创建
                if not os.path.exists(train_dir):
                    os.mkdir(train_dir)
                if not os.path.exists(val_dir):
                    os.mkdir(val_dir)
                if not os.path.exists(test_dir):
                    os.mkdir(test_dir)
                # 处理每个视频
                for video in train:
                    self.process_video(video, file, train_dir)
    
                for video in val:
                    self.process_video(video, file, val_dir)
    
                for video in test:
                    self.process_video(video, file, test_dir)
    
            print('Preprocessing finished.')
    
        def process_video(self, video, action_name, save_dir):
            video_filename = video.split('.')[0]# 视频文件名，在这是文件夹
            # 不存在即创建
            if not os.path.exists(os.path.join(save_dir, video_filename)):
                os.mkdir(os.path.join(save_dir, video_filename))
            # 读取每个视频
            capture = cv2.VideoCapture(os.path.join(self.root_dir, action_name, video))
            # 得到视频的帧数和宽高
            frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
            frame_width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
            frame_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
            # 确保都能得到16帧图像，不满足的就每隔3、2、1帧取一次图像
            EXTRACT_FREQUENCY = 4 # 每隔4帧取一次图像
            if frame_count // EXTRACT_FREQUENCY <= 16:
                EXTRACT_FREQUENCY -= 1
                if frame_count // EXTRACT_FREQUENCY <= 16:
                    EXTRACT_FREQUENCY -= 1
                    if frame_count // EXTRACT_FREQUENCY <= 16:
                        EXTRACT_FREQUENCY -= 1
    
            count = 0
            i = 0
            retaining = True
    
            while (count < frame_count and retaining):
                retaining, frame = capture.read()# 读每一帧图像（capture.read():读没读出来，读出来的数据）
                if frame is None:
                    continue
    
                if count % EXTRACT_FREQUENCY == 0:# 每隔4帧取一次图像
                    # resize图像
                    if (frame_height != self.resize_height) or (frame_width != self.resize_width):
                        frame = cv2.resize(frame, (self.resize_width, self.resize_height))
                    # 写入图像
                    cv2.imwrite(filename=os.path.join(save_dir, video_filename, '0000{}.jpg'.format(str(i))), img=frame)
                    i += 1
                count += 1
            # 释放视频捕获
            capture.release()
    
        def randomflip(self, buffer):
            """Horizontally flip the given image and ground truth randomly with a probability of 0.5."""
    
            if np.random.random() < 0.5:
                for i, frame in enumerate(buffer):
                    frame = cv2.flip(buffer[i], flipCode=1)
                    buffer[i] = cv2.flip(frame, flipCode=1)
    
            return buffer
    
    
        def normalize(self, buffer):
            for i, frame in enumerate(buffer):
                frame -= np.array([[[90.0, 98.0, 102.0]]])
                buffer[i] = frame
    
            return buffer
    
        def to_tensor(self, buffer):# 网络的输入c×l×h×w（通道×帧数×高×宽）
            return buffer.transpose((3, 0, 1, 2))
    
        def load_frames(self, file_dir):
            frames = sorted([os.path.join(file_dir, img) for img in os.listdir(file_dir)])# 拿到每张图像
            frame_count = len(frames)# 计算图像序列的长度，不是16的要筛选掉
            # 做缓存,分配了一块未被清空的内存
            buffer = np.empty((frame_count, self.resize_height, self.resize_width, 3), np.dtype('float32'))
            # 读数据
            for i, frame_name in enumerate(frames):
                frame = np.array(cv2.imread(frame_name)).astype(np.float64)
                buffer[i] = frame
    
            return buffer
    
        def crop(self, buffer, clip_len, crop_size):# 规则：从一段到另一段一共取16帧
            # 随机选择时间索引。要取帧数开始的位置
            time_index = np.random.randint(buffer.shape[0] - clip_len)
            # 切割图像到112x112，还是取索引的操作
            height_index = np.random.randint(buffer.shape[1] - crop_size)
            width_index = np.random.randint(buffer.shape[2] - crop_size)
    
            buffer = buffer[time_index:time_index + clip_len,
                     height_index:height_index + crop_size,
                     width_index:width_index + crop_size, :]
    
            return buffer
    
    if __name__ == "__main__":
        from torch.utils.data import DataLoader
        train_data = VideoDataset(dataset='ucf101', split='train', clip_len=16, preprocess=False)
        train_loader = DataLoader(train_data, batch_size=100, shuffle=True)
    
        for i, sample in enumerate(train_loader):
            inputs = sample[0]
            labels = sample[1]
            print(inputs.size())
            print(labels)
    
            if i == 1:
                break

做好的效果：  
![在这里插入图片描述][20200720170051417.png]  
![在这里插入图片描述][20200720170110184.png]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhMXR6eQ_size_16_color_FFFFFF_t_70]

每个test、train、val文件夹下101个文件夹，为类别名(标签)。

[20200720170051417.png]: /images/20230209/13026bae287149558bbd05f7f363657b.png
[20200720170110184.png]: https://img-blog.csdnimg.cn/20200720170110184.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhMXR6eQ_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/2020072018540998.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dhMXR6eQ==,size_16,color_FFFFFF,t_70