Spaces:
Sleeping
Sleeping
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) |