Python视频去背景

Python视频去背景

目前需求：需要将视频的背景去掉，只保留人像，方便后期合成其他视频

大概步骤：

1、先对视频每一帧进行图像分割

2、对每一帧的图像进行抠图

3.对完成抠图的每一帧图像进行合并

使用之前先导入库

cv2:进行图片分割

PIL：操作图像

os：操作系统目录

rembg ：对图像进行抠图移除背景

numpy：图片分割时使用

import cv2
from PIL import Image
import os
from rembg import remove
import numpy as np

分割视频为图像的代码

# 读取视频文件
video_path = 'D:\\vidio\\1.mp4'
cap = cv2.VideoCapture(video_path)
# 获取视频帧率和总帧数
fps = cap.get(cv2.CAP_PROP_FPS)
# 创建保存照片的目录
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if not os.path.exists('frames'):
    os.makedirs('frames')
# 分割视频为照片
for i in range(total_frames):
    ret, frame = cap.read()
    if ret:
        # 每秒25帧
        if i % int(fps/25) == 0:
            # 保存照片
            cv2.imwrite(f'frames/frame_{i}.jpg', frame)
    else:
        break

移除图片背景代码

# 处理照片
for filename in os.listdir('frames'):
    if filename.endswith('.jpg'):
        # 打开照片
        img = Image.open(f'frames/{filename}')
        # 抠图
        img = remove(img)
        # 保存处理后的照片
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGR)
        # 保存处理后的照片
        cv2.imwrite(f'processed_frames/{filename}', img)

合并图片为视频代码

# 设置输入图像目录和输出视频文件名
input_dir = "D:\\pythonProject1\\processed_frames"
output_file = "output_video.mp4"

# 获取输入目录中的所有图像文件
image_files = [f for f in os.listdir(input_dir) if f.endswith(".jpg")]

# 从第一张图像获取图像的尺寸
img = cv2.imread(os.path.join(input_dir, image_files[0]))
height, width, _ = img.shape

# 创建用于写入视频的VideoWriter对象
video_writer = cv2.VideoWriter(output_file, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height))

# 逐个读取图像并将其写入视频
for image_file in image_files:
    img = cv2.imread(os.path.join(input_dir, image_file))
    video_writer.write(img)

所有代码

import cv2
from PIL import Image
import os
from rembg import remove
import numpy as np
# 读取视频文件
video_path = 'D:\\vidio\\1.mp4'
cap = cv2.VideoCapture(video_path)
# 获取视频帧率和总帧数
fps = cap.get(cv2.CAP_PROP_FPS)
# 创建保存照片的目录
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if not os.path.exists('frames'):
    os.makedirs('frames')
# 分割视频为照片
for i in range(total_frames):
    ret, frame = cap.read()
    if ret:
        # 每秒25帧
        if i % int(fps/25) == 0:
            # 保存照片
            cv2.imwrite(f'frames/frame_{i}.jpg', frame)
    else:
        break

# 处理照片
for filename in os.listdir('frames'):
    if filename.endswith('.jpg'):
        # 打开照片
        img = Image.open(f'frames/{filename}')
        # 抠图
        img = remove(img)
        # 保存处理后的照片
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGR)
        # 保存处理后的照片
        cv2.imwrite(f'processed_frames/{filename}', img)

# 设置输入图像目录和输出视频文件名
input_dir = "D:\\pythonProject1\\processed_frames"
output_file = "output_video.mp4"

# 获取输入目录中的所有图像文件
image_files = [f for f in os.listdir(input_dir) if f.endswith(".jpg")]

# 从第一张图像获取图像的尺寸
img = cv2.imread(os.path.join(input_dir, image_files[0]))
height, width, _ = img.shape

# 创建用于写入视频的VideoWriter对象
video_writer = cv2.VideoWriter(output_file, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height))

# 逐个读取图像并将其写入视频
for image_file in image_files:
    img = cv2.imread(os.path.join(input_dir, image_file))
    video_writer.write(img)

# 释放VideoWriter对象并关闭文件
video_writer.release()
cv2.destroyAllWindows()