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 .matting_model.model import MattingNetwork
|
import uuid
|
|
|
VARIANT_RESNET50 ="resnet50"
|
CHECKPOINT_RESNET50 ="./models/rvm_resnet50.pth"
|
|
VARIANT_MOBILENETV3 ="mobilenetv3"
|
CHECKPOINT_MOBILENETV3 ="./models/rvm_mobilenetv3.pth"
|
|
|
class VideoMattingConverter:
|
"""
|
视频抠像转换器
|
封装原来的 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_jpg: 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_jpg and not os.path.exists(bg_jpg):
|
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_jpg,
|
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, int(output_video_mbps * 1_000_000),input_source)
|
else:
|
writer_com = VideoWriter(output_composition, frame_rate, int(output_video_mbps * 1_000_000))
|
if output_alpha:
|
writer_pha = VideoWriter(output_alpha, frame_rate, int(output_video_mbps * 1_000_000))
|
if output_foreground:
|
writer_fgr = VideoWriter(output_foreground, frame_rate, 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:
|
bgr = torch.tensor([155, 255, 155], 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
|
