Spaces:
Sleeping
Sleeping
File size: 3,746 Bytes
b0c0f32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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) |