CelebDF_dataset

import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import cv2
from sklearn.model_selection import train_test_split


class CelebDFDataset(Dataset):

    def __init__(self, root_dir, output_dir, label_file_path, split='train', clip_len=16, preprocess=False):
        self.root_dir = root_dir
        self.output_dir = output_dir
        self.folder = os.path.join(root_dir, split)
        self.videos = os.listdir(self.folder)
        self.label_file_full_path = os.path.join(root_dir, label_file_path)
        self.video_label_map = {}  # {id0_0000.mp4:1, id0_id1_0000:0}
        self.clip_len = clip_len
        self.split = split

        # 图片resize参数
        self.resize_height = 128
        self.resize_width = 171
        self.crop_size = 112

        if (not self.check_preprocess()) or preprocess:
            print("Preprocessing of dataset, this will take long, but it will be done only once.")
            self.preprocess()

        self.fnames, labels = [], []
        # folder：大类文件夹
        # frames: 每一个视频对应的文件夹，存储所有数据，标签等于其上层目录的label
        for label in sorted(os.listdir(self.folder)):
            for fname in os.listdir(os.path.join(self.folder, label)):
                self.fnames.append(os.path.join(self.folder, label, fname))
                labels.append(label)

        assert len(labels) == len(self.fnames)
        print("Number of {} videos:{:d}".format(split, len(self.fnames)))

        # 创建字符串类型的标签名字与索引的对应
        self.label2index = {label: index for index, label in enumerate(sorted(set(labels)))}
        # 将string标签转化为int标签
        self.label_array = np.array([self.label2index[label] for label in labels], dtype=int)

    def __len__(self):
        return len(os.listdir(self.root_dir))

    def __getitem__(self, index):
        # 读取对应某一个视频的所有视频帧，预处理
        buffer = self.load_frames(self.fnames[index])
        buffer = self.crop(buffer, self.clip_len, self.crop_size)
        labels = np.array(self.label_array[index])

        if self.split == 'test':
            buffer = self.randomflip(buffer)
        buffer = self.normalize(buffer)
        buffer = self.to_tensor(buffer)
        return torch.from_numpy(buffer), torch.from_numpy(labels)

    def _build_label_file_map(self):
        with open(self.label_file_full_path, 'r') as f:
            label_file_list = f.readlines()
            for label_file in label_file_list:
                label_file = label_file.replace("\n", "")
                label = int(label_file.split("\t")[0])
                file = label_file.split("\t")[1]
                self.video_label_map[file] = label

    def check_preprocess(self):
        pass

    def preprocess(self):
        # 创建输出结果子文件夹
        print("output_dir is" + str(self.output_dir))
        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'))

        # 划分train/val/test sets
        for file in os.listdir(self.root_dir):
            file_path = os.path.join(self.root_dir, file)  # file表示每一个视频文件夹
            video_files = [name for name in os.listdir(file_path)]  # 每一类视频中的视频文件
            # train/val/test划分比例为0.64：0.16：0.2
            train_and_valid, test = train_test_split(video_files, test_size=0.2, random_state=42)
            train, val = train_and_valid(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, fake, 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, fake, 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帧，默认抽帧频率为4，最低为1，如果视频总长度少于16帧，会报错
        EXTRACT_FREQUENCY = 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()
            if frame is None:
                continue

            if count % EXTRACT_FREQUENCY == 0:
                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):

    def normalize(self, buffer):

    def to_tensor(self, buffer):
        for frame in buffer:
            

    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)
        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):
        if buffer.shape[0] <= clip_len:
            print("该视频没有足够的帧数可供选择")
            time_index = 0
        else:
            time_index = np.random.randint(buffer.shape[0] - clip_len)

        # 随机选择空间裁剪参数
        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