整理了一下大概步骤如下：
1、生成样本，分别从有码的视频和无码的视频抽取1000张或者更多视频截图，越多训练出来的模型越准确；
2、开始训练样本，训练好的模型叫mosaic_detector.h5，8.95 MB
3、模型验证
4、用训练好的模型来分拣视频
下面是4个步骤的源码备份：
1、保存视频截图.py
import cv2
import os
import random
# 定义全局变量
INTERVAL = 1000 # 每隔多少帧抽取一张图片
def extract_frames_single_thread(video_folder, output_folder, label):
"""
使用单线程从指定文件夹中的所有视频中提取帧并保存到指定文件夹。
参数：
- video_folder: 视频文件夹路径。
- output_folder: 输出图像的保存文件夹。
- label: 'mosaic' 或 'clean'，用于区分马赛克和非马赛克图像。
"""
video_files = [os.path.join(video_folder, f) for f in os.listdir(video_folder)
if f.lower().endswith(('.mp4', '.avi', '.mkv', '.mov', '.wmv'))]
if not video_files:
print("未找到符合条件的视频文件。")
return
else:
print("待处理的视频如下---------------------------：")
_ = [print(i) for i in video_files]
if not os.path.exists(output_folder):
os.makedirs(output_folder)
for video_path in video_files:
print('正在处理----------------------------------------------------------',video_path)
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
save_count = 0
if total_frames == 0:
print(f"无法读取视频文件或视频文件为空: {video_path}")
continue
for frame_count in range(0, total_frames, INTERVAL):
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
ret, frame = cap.read()
if not ret:
print(f"无法读取帧或已达到视频末尾: {video_path}")
break
random_number = random.randint(100000, 999999) # 生成一个6位随机数
filename = f"{label}_{frame_count}_{random_number}.jpg"
filepath = os.path.join(output_folder, filename)
cv2.imwrite(filepath, frame)
save_count += 1
cap.release()
print(f"从 {video_path} 提取了 {save_count} 张图像。")
# 示例用法
# 提取指定文件夹中的所有视频帧
video_folder = 'H:/迅雷云盘/'
output_folder = 'data/mosaic'
label = 'mosaic'
extract_frames_single_thread(video_folder, output_folder, label)
2、开始训练.py
import os
import numpy as np
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
from tensorflow.keras.preprocessing.image import load_img, img_to_array
# 获取图像的原始尺寸
def get_image_size(image_path):
img = load_img(image_path)
return img.size
# 数据目录
train_data_dir = 'data'
validation_data_dir = 'data'
# 获取训练集和验证集中任意一张图像的尺寸
sample_image_path = os.path.join(train_data_dir, 'mosaic', os.listdir(os.path.join(train_data_dir, 'mosaic'))[0])
img_width, img_height = get_image_size(sample_image_path)
print(f"使用原始图像尺寸：{img_width}x{img_height}")
# 检查数据集平衡性
mosaic_count = len(os.listdir(os.path.join(train_data_dir, 'mosaic')))
clean_count = len(os.listdir(os.path.join(train_data_dir, 'clean')))
print(f"马赛克图像数: {mosaic_count}, 清晰图像数: {clean_count}")
# 数据生成器与数据增强
train_datagen = ImageDataGenerator(
rescale=1./255,
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
horizontal_flip=True,
validation_split=0.2 # 20% 用于验证
)
# 调整后的批量大小批量大小（Batch Size）批量大小影响每个周期处理的数据量。较大的批量大小通常会加快训练速度，但会增加内存占用。调整批量大小可以影响训练速度。
batch_size = 16 # 从32减少到16
# 训练集生成器
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size=(img_width, img_height),
batch_size=batch_size,
class_mode='binary',
subset='training'
)
# 验证集生成器
validation_generator = train_datagen.flow_from_directory(
validation_data_dir,
target_size=(img_width, img_height),
batch_size=batch_size,
class_mode='binary',
subset='validation'
)
# 加载预训练模型（不包括顶层）
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(img_width, img_height, 3))
# 冻结预训练模型的层
base_model.trainable = False
# 构建模型
model = Sequential([
base_model,
GlobalAveragePooling2D(),
Dropout(0.5),
Dense(1, activation='sigmoid')
])
# 调整学习率并编译模型
model.compile(optimizer=Adam(lr=0.00001), loss='binary_crossentropy', metrics=['accuracy'])
# 回调函数
checkpoint = ModelCheckpoint('mosaic_detector.h5', monitor='val_accuracy', verbose=1, save_best_only=True, mode='max')
earlystop = EarlyStopping(monitor='val_accuracy', patience=5, verbose=1, mode='max')
callbacks_list = [checkpoint, earlystop]
# 训练模型
model.fit(
train_generator,
steps_per_epoch=train_generator.samples // batch_size,
epochs=10,
validation_data=validation_generator,
validation_steps=validation_generator.samples // batch_size,
callbacks=callbacks_list
)
3、验证模型准确率.py
import os
import time # 确保导入 time 模块
import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
import shutil
import multiprocessing as mp
import tensorflow as tf
# 设置环境变量以禁用 CPU 特性提示
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
# 禁用 TensorFlow 的日志输出
tf.get_logger().setLevel('ERROR')
# 加载训练好的模型
model = load_model('mosaic_detector.h5')
# 数据目录
test_data_dir = 'data/test'
you_data_dir = os.path.join(test_data_dir, 'you')
# 如果不存在，创建 you 文件夹
if not os.path.exists(you_data_dir):
os.makedirs(you_data_dir)
# 准备预测数据
test_images = [f for f in os.listdir(test_data_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
total_images = len(test_images)
# 单个文件的处理函数
def process_image(img_name):
img_path = os.path.join(test_data_dir, img_name)
img = image.load_img(img_path, target_size=(1280, 720)) # 使用模型训练时的输入尺寸
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_array /= 255.0
prediction = model.predict(img_array, verbose=0)
if prediction > 0.34:
# 预测为马赛克图像，移动到 you 文件夹
shutil.move(img_path, os.path.join(you_data_dir, img_name))
return f"{img_name} {prediction} 包含马赛克，已移动到 {you_data_dir}"
else:
return f"{img_name} {prediction} 未包含马赛克，保持不变"
# 多进程处理函数
def process_images_in_parallel():
with mp.Pool(processes=1) as pool:
for img_index, result in enumerate(pool.imap(process_image, test_images), start=1):
print(f"{img_index}/{total_images}: {result}")
if __name__ == '__main__':
# 统计总处理时间
start_time = time.time()
# 执行多进程处理
process_images_in_parallel()
print(f"有马赛克的图像已移动到 {you_data_dir} 文件夹中")
# 统计总处理时间
end_time = time.time()
total_time = end_time - start_time
print(f"总处理时间: {total_time:.2f} 秒")
4、用来分拣视频.py
import os
import time
import shutil
import cv2
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
# 加载保存的整个模型，保持原始输入形状
model = load_model('mosaic_detector.h5')
# 常见的视频格式（不分大小写）
video_formats = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.mpg', '.mpeg', '.3gp', '.webm']
###数据目录
source_dir = 'H:/迅雷云盘/' # 需要处理的视频目录
destination_dir = os.path.join(source_dir, 'youma') # 含马赛克视频的目标目录
# 如果不存在，创建目标目录
if not os.path.exists(destination_dir):
os.makedirs(destination_dir)
# 检查视频是否包含马赛克
def contains_mosaic(video_path, model, input_shape):
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
step = 1000 # 每隔1000帧检测一次
print(f"正在处理文件: {video_path}")
print(f"总帧数: {total_frames}")
for frame_count in tqdm(range(0, total_frames, step), desc="处理进度"):
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
ret, frame = cap.read()
if not ret:
continue
# 调整帧的尺寸以匹配模型的输入形状
frame = cv2.resize(frame, (input_shape[1], input_shape[0])) # 调整图像尺寸为模型的输入形状
frame_array = img_to_array(frame)
frame_array = np.expand_dims(frame_array, axis=0)
frame_array /= 255.0
prediction = model.predict(frame_array)
if prediction > 0.5: # 使用阈值0.5检测
cap.release()
return True
cap.release()
return False
# 处理单个文件
def process_file(file):
video_path = file['path']
model = file['model']
input_shape = file['input_shape']
if contains_mosaic(video_path, model, input_shape):
shutil.move(video_path, os.path.join(destination_dir, os.path.basename(video_path)))
result = f"{os.path.basename(video_path)} 包含马赛克，已移动到 {destination_dir}"
else:
result = f"{os.path.basename(video_path)} 未包含马赛克，保持不变"
return result
# 获取要检测的视频文件
video_files = []
input_shape = (1280, 720, 3) # 模型期望的输入形状
for root, dirs, files in os.walk(source_dir):
for file in files:
if any(file.lower().endswith(ext) for ext in video_formats):
video_path = os.path.join(root, file)
video_files.append({'path': video_path, 'model': model, 'input_shape': input_shape})
# 统计总处理时间
start_time = time.time()
# 使用多线程处理文件
with ThreadPoolExecutor(max_workers=4) as executor: # 你可以调整最大线程数
future_to_file = {executor.submit(process_file, file): file for file in video_files}
for future in tqdm(as_completed(future_to_file), total=len(video_files), desc="总处理进度"):
print(future.result())
# 统计总处理时间
end_time = time.time()
total_time = end_time - start_time
print(f"总处理时间: {total_time:.2f} 秒")
# import os
# import time
# import shutil
# import cv2
# import numpy as np
# from tqdm import tqdm
# from tensorflow.keras.models import load_model
# from tensorflow.keras.preprocessing.image import img_to_array
# # 加载保存的整个模型，保持原始输入形状
# model = load_model('mosaic_detector.h5')
# # 常见的视频格式（不分大小写）
# video_formats = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.mpg', '.mpeg', '.3gp', '.webm']
# # 数据目录
# source_dir = 'H:/迅雷云盘/' # 需要处理的视频目录
# destination_dir = os.path.join(source_dir, 'youma') # 含马赛克视频的目标目录
# # 如果不存在，创建目标目录
# if not os.path.exists(destination_dir):
# os.makedirs(destination_dir)
# # 检查视频是否包含马赛克
# def contains_mosaic(video_path, model, input_shape):
# cap = cv2.VideoCapture(video_path)
# total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# step = 1000 # 每隔1000帧检测一次
# for frame_count in tqdm(range(0, total_frames, step), desc=f"处理进度: {os.path.basename(video_path)}"):
# cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
# ret, frame = cap.read()
# if not ret:
# continue
# # 调整帧的尺寸以匹配模型的输入形状
# frame = cv2.resize(frame, (input_shape[1], input_shape[0])) # 调整图像尺寸为模型的输入形状
# frame_array = img_to_array(frame)
# frame_array = np.expand_dims(frame_array, axis=0)
# frame_array /= 255.0
# prediction = model.predict(frame_array)
# if prediction > 0.5: # 使用阈值0.5检测
# cap.release()
# return True
# cap.release()
# return False
# # 统计总处理时间
# start_time = time.time()
# # 获取要检测的视频文件
# video_files = []
# input_shape = (1280, 720, 3) # 模型期望的输入形状
# for root, dirs, files in os.walk(source_dir):
# for file in files:
# if any(file.lower().endswith(ext) for ext in video_formats):
# video_path = os.path.join(root, file)
# video_files.append(video_path)
# # 单线程处理文件
# total_videos = len(video_files)
# for index, video_path in enumerate(video_files, start=1):
# tqdm.write(f"总处理进度: {index}/{total_videos} - 正在处理文件: {video_path} (总帧数: {cv2.VideoCapture(video_path).get(cv2.CAP_PROP_FRAME_COUNT):.0f})")
# if contains_mosaic(video_path, model, input_shape):
# shutil.move(video_path, os.path.join(destination_dir, os.path.basename(video_path)))
# tqdm.write(f"{os.path.basename(video_path)} 包含马赛克，已移动到 {destination_dir}")
# else:
# tqdm.write(f"{os.path.basename(video_path)} 未包含马赛克，保持不变")
# # 更新进度条描述
# tqdm.write(f"总处理进度: {index}/{total_videos}")
# # 统计总处理时间
# end_time = time.time()
# total_time = end_time - start_time
# print(f"总处理时间: {total_time:.2f} 秒")
评论