import cv2
import torch
import torchvision.models as models
import torchvision.transforms as transforms
import numpy as np
import pickle
from PIL import Image
import os


def load_model_vgg19():

    # 加载预训练的 VGG19 模型
    vgg19 = models.vgg19(pretrained=True)


    # structure = torch.nn.Sequential(*list(vgg19.children())[:])
    # # 查看整体结构
    # print(structure) 
    # # 查看模型各部分名称  ['features', 'avgpool', 'classifier']
    # print('模型各部分名称', vgg19._modules.keys())
    # # 原始feature
    # features = torch.nn.Sequential(*list(vgg19.children())[0])
    # # 原始classifier
    # classifier = torch.nn.Sequential(*list(vgg19.children())[-1])

    new_classifier = torch.nn.Sequential(*list(vgg19.children())[-1][:4])

    vgg19.classifier = new_classifier
    vgg19.eval()

    return vgg19


# 定义图像预处理步骤
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 处理视频，提取帧并特征，并保存为 .pkl 文件
def extract_frames_and_features(model_vgg19, video_path):
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    feature_list = []
    frame_count = 0
    buffer = []

    count = 0
    while True:
        # 读取视频帧
        ret, frame = cap.read()
        if not ret:
            break
        # 将 BGR 格式的帧转换为 RGB 格式
        frame_count += 1
        if frame_count % 16 == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            # 将帧转换为 PIL.Image
            pil_frame = Image.fromarray(frame_rgb)
            # 将帧进行预处理
            input_tensor = transform(pil_frame)
            # 添加一个批次维度
            input_batch = input_tensor.unsqueeze(0)
            # 使用 VGG19 特征提取器提取特征
            with torch.no_grad():
                output = model_vgg19(input_batch).data[0]
                feature_list.append(output.numpy()) 
    # 关闭视频文件
    cap.release()
    feature_matrix = np.stack(feature_list)
    print(feature_matrix.shape)
    return feature_matrix

# # 视频文件路径
# video_path = "/mnt/data10t/dazuoye/GROUP2024-GEN6/FakeSV/code/C3D_Feature_Extractor/raw_video/douyin_6571001202379590925.mp4"
# # 提取视频帧并特征
# video_features = extract_frames_and_features(video_path)
# # 将特征列表保存为 .pkl 文件
# # 确保输出目录存在，如果不存在则创建
# output_dir = 'outputs'
# if not os.path.exists(output_dir):
#     os.makedirs(output_dir)
# output_file_path = os.path.join(output_dir, 'video_features.pkl')
# with open(output_file_path, 'wb') as f:
#     pickle.dump(video_features, f)
# print("Video features saved to", output_file_path)

def process_video(model_vgg19, video_name, input_folder, output_folder):
    # 确保输出目录存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    video_path = os.path.join(input_folder, video_name)
    video_features = extract_frames_and_features(model_vgg19, video_path)
    output_file_path = os.path.join(output_folder, video_name.split('.')[0] + '.pkl')
    with open(output_file_path, 'wb') as f:
        pickle.dump(video_features, f)
    print("Video features saved to", output_file_path)

if __name__ == "__main__":
    input_folder = './FakeVD/code/C3D_Feature_Extractor/raw_video'
    output_folder = './FakeVD/code/VGG19/outputs'
    video_name = 'douyin_6571001202379590925.mp4'

    model_vgg19 = load_model_vgg19()
    process_video(model_vgg19, video_name, input_folder, output_folder)