import os import torch from torch.utils.data import DataLoader from torchvision import transforms from typing import Optional, Tuple from PIL import Image import time from pathlib import Path from auto_media_publisher.utils.media_input_output import VideoReader, VideoWriter, ImageSequenceReader, ImageSequenceWriter from .robustvideomatting.model import MattingNetwork import uuid VARIANT_RESNET50 ="resnet50" CHECKPOINT_RESNET50 ="./models/rvm_resnet50.pth" VARIANT_MOBILENETV3 ="mobilenetv3" CHECKPOINT_MOBILENETV3 ="./models/rvm_mobilenetv3.pth" class RobustVideoMattingConverter: """ 视频抠像转换器 封装原来的 convert_video 逻辑，支持 video / png_sequence 输出 """ def __init__(self, precision:str = 'high', device: str = "cuda"): self.device = device if precision=='high': self.model = MattingNetwork(VARIANT_RESNET50).eval().to(device) self.model.load_state_dict(torch.load(CHECKPOINT_RESNET50, map_location=device)) else: self.model = MattingNetwork(VARIANT_MOBILENETV3).eval().to(device) self.model.load_state_dict(torch.load(CHECKPOINT_MOBILENETV3, map_location=device)) self.model = torch.jit.script(self.model) self.model = torch.jit.freeze(self.model) @staticmethod def auto_downsample_ratio(h: int, w: int) -> float: """自动计算 downsample_ratio""" return min(512 / max(h, w), 1) def video_matting(self, input: str, bg_video: Optional[str] = None, bg_image: Optional[str] = None, output_video: Optional[str] = None, output_alpha: Optional[str] = None, output_fore: Optional[str] = None): assert any([output_video, output_alpha, output_fore]), 'Must provide at least one output.' if not os.path.exists(input): print("input file is NOT exist.") return False if bg_video and not os.path.exists(bg_video): print("input background video is NOT exist.") return False if bg_image and not os.path.exists(bg_image): print("input background image is NOT exist.") return False return self._video_matting_core( input_source = input, input_resize = None, downsample_ratio = None, output_type = 'video', background_video = bg_video, background_jpg = bg_image, with_audio = True, output_composition = output_video, output_alpha = output_alpha, output_foreground = output_fore, output_video_mbps = None, dtype =torch.float32) def _video_matting_core(self, input_source: str, input_resize: Optional[Tuple[int, int]] = None, downsample_ratio: Optional[float] = None, output_type: str = 'video', background_video: Optional[str] = None, background_jpg: Optional[str] = None, with_audio: bool = True, output_composition: Optional[str] = None, output_alpha: Optional[str] = None, output_foreground: Optional[str] = None, output_video_mbps: Optional[float] = None, dtype: Optional[torch.dtype] = torch.float32): assert downsample_ratio is None or (0 < downsample_ratio <= 1), 'Downsample ratio must be between 0 and 1.' assert any([output_composition, output_alpha, output_foreground]), 'Must provide at least one output.' assert output_type in ['video', 'png_sequence'], 'Only support video or png_sequence.' success = False # 初始化 transform if input_resize is not None: transform = transforms.Compose([ transforms.Resize(input_resize[::-1]), transforms.ToTensor() ]) else: transform = transforms.ToTensor() # 读取输入源 if os.path.isfile(input_source): source = VideoReader(input_source, transform) else: source = ImageSequenceReader(input_source, transform) reader = DataLoader(source, batch_size=1, pin_memory=True, num_workers=0) # 初始化输出 if output_type == 'video': frame_rate = source.frame_rate if isinstance(source, VideoReader) else 30 output_video_mbps = 1 if output_video_mbps is None else output_video_mbps if output_composition: if with_audio: writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000),input_source) else: writer_com = VideoWriter(output_composition, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000)) if output_alpha: writer_pha = VideoWriter(output_alpha, frame_rate, source.width,source.height,int(output_video_mbps * 1_000_000)) if output_foreground: writer_fgr = VideoWriter(output_foreground, frame_rate,source.width,source.height, int(output_video_mbps * 1_000_000)) else: # png_sequence if output_composition: writer_com = ImageSequenceWriter(output_composition, 'png') if output_alpha: writer_pha = ImageSequenceWriter(output_alpha, 'png') if output_foreground: writer_fgr = ImageSequenceWriter(output_foreground, 'png') # 背景设置 background_video_reader = None if output_composition and output_type == 'video': if background_video: background_video_reader = VideoReader(background_video, transform=transforms.ToTensor()) elif background_jpg: bgr_IMG = Image.open(background_jpg) bgr = transforms.ToTensor()(bgr_IMG).to(self.device) else: # 使用 0,200,0 而不是 0,255,0 避免过度曝光 bgr = torch.tensor([0, 200, 0], device=self.device, dtype=dtype).div(255).view(1, 1, 3, 1, 1) # 开始推理 self.model = self.model.eval() try: with torch.no_grad(): rec = [None] * 4 current_bg_idx = 1 bg_increasing = True for src in reader: if downsample_ratio is None: downsample_ratio = self.auto_downsample_ratio(*src.shape[2:]) src = src.to(self.device, dtype, non_blocking=True).unsqueeze(0) fgr, pha, *rec = self.model(src, *rec, downsample_ratio) # 前景 if output_foreground: writer_fgr.write(fgr[0]) if output_alpha: writer_pha.write(pha[0]) # 合成图 if output_composition: if output_type == 'video': if background_video: tmp_bgr = background_video_reader[current_bg_idx].to(self.device) if bg_increasing: current_bg_idx += 1 if current_bg_idx >= len(background_video_reader): current_bg_idx = len(background_video_reader) - 2 bg_increasing = False else: current_bg_idx -= 1 if current_bg_idx < 1: current_bg_idx = 1 bg_increasing = True com = fgr * pha + tmp_bgr * (1 - pha) else: com = fgr * pha + bgr * (1 - pha) else: fgr = fgr * pha.gt(0) com = torch.cat([fgr, pha], dim=-3) writer_com.write(com[0]) # process finished. success = True finally: # Clean up if output_composition is not None: writer_com.close() ''' time.sleep(0.5) # 确保文件释放，Windows 上非常必要 if success and with_audio and os.path.exists(output_composition) : out_dir = os.path.dirname(output_composition) base_name = os.path.basename(output_composition) audio_output_path = os.path.join(out_dir, f"audio_{base_name}") os.system(f'ffmpeg -i {output_composition} -i {input_source} -c copy ' f'-map 0:v:0 -map 1:a:0 {audio_output_path}') ''' if output_alpha: writer_pha.close() if output_foreground: writer_fgr.close() return success