保存带有透明通道的视频:
import osimport imageio
from rembg import remove as removBg,new_session
from PIL import Image
import numpy as np
import cv2
from tqdm import tqdmclass cls_rembg():def __init__(self,model_pth):self.session = new_session(model_pth)# 替换img_src 背景图路径为img_back_path,img_back_path == None 为透明# 输入图片为 pil image 输出为替换好的 pil imagedef remove_background(self,img_src,bgcolor = None,img_back_path = None,needRemoveBg = True):imgret_newbg = img_srcoutput_png = img_srcif needRemoveBg:if img_back_path is not None:output_png = removBg(data = img_src,session=self.session,alpha_matting_erode_size = 5)imgret_newbg = self.replace_background(output_png,img_back_path=img_back_path)elif bgcolor is not None:output_png = removBg(data = img_src,session=self.session,bgcolor = bgcolor,alpha_matting_erode_size = 5)imgret_nobg = output_pngreturn imgret_newbg,imgret_nobgif __name__ == '__main__':rembg = cls_rembg('u2net_human_seg.onnx')video_stream = cv2.VideoCapture("1.mp4")face_nobg = 'test_nobg.mov'#face_nobg = 'test_nobg.webm'face_newbg = 'test_newbg.mp4'video_writer = imageio.get_writer(face_nobg, format='FFMPEG', mode='I', codec='png', fps=30,ffmpeg_params=['-pix_fmt', 'rgba'])#video_writer = imageio.get_writer(face_nobg, format='FFMPEG', mode='I', codec='libvpx-vp9', fps=30)frames_nobg = []frames_newbg = []full_frames = []fps = video_stream.get(cv2.CAP_PROP_FPS)while True:still_reading, frame = video_stream.read()if not still_reading:video_stream.release()breakfull_frames.append(frame)back_color = (255, 255, 255, 0)if len(full_frames) > 0:for idx in tqdm(range(len(full_frames)), 'remove/replace background'):frame_face = Image.fromarray(cv2.cvtColor(full_frames[idx], cv2.COLOR_BGRA2RGBA))frame_newbg, frame_nobg = rembg.remove_background(frame_face, back_color, None)frames_nobg.append(frame_nobg)image = np.array(frame_nobg)# 确保图像的尺寸与第一个图像相同#if image.shape[:2] != (700, 700):#raise ValueError("All images should have the same dimensions.")# 将图像的 alpha 通道分离出来#alpha_channel = image[:, :, 3]# 将图像的 alpha 通道与 RGB 通道合并,得到 RGBA 图像#rgba_image = np.dstack((image[:, :, :3], alpha_channel))video_writer.append_data(image)video_writer.close()#if len(frames_nobg) > 0:#frames_nobg[0].save(face_nobg, save_all=True, append_images=frames_nobg, fps=float(fps))
保存带有透明通道的视频:
from moviepy.editor import VideoFileClip, CompositeVideoClip
from moviepy.video.io.ffmpeg_reader import FFMPEG_VideoReader
import imageiozm_video_path = "test_nobg.mov"if __name__ == '__main__':#video_stream = imageio.get_reader("test_nobg.mov", format='FFMPEG', fps=30,# ffmpeg_params=['-pix_fmt', 'rgba'])video_stream = VideoFileClip('test_nobg.mov', has_mask=True).readerframes_nobg = []frames_newbg = []full_frames = []for num, frame in enumerate(video_stream.read_frame()):full_frames.append(frame)print('{}'.format(len(full_frames)))