from auto_media_publisher.modifiers.video_matting.matanyone.inference.inference_core import InferenceCore
|
from auto_media_publisher.utils.utils_yolo import generate_mask_from_first_frame
|
from auto_media_publisher.utils.logger import get_logger
|
import os
|
import traceback
|
MAT_ANYONE_MODEL_DIR ="./models/MatAnyone"
|
|
|
class MatAnyoneConverter:
|
"""
|
视频抠像转换器-- 只转换多人中的一人。
|
封装原来的
|
"""
|
|
def __init__(self):
|
self.logger = get_logger('MatAnyoneConverter','admin')
|
self.processor = InferenceCore(MAT_ANYONE_MODEL_DIR)
|
|
|
def video_matting(self,input_path,output_path,bg_video=None,bg_image=None,filter_type="dist"):
|
if not input_path or (input_path and not os.path.exists(input_path)):
|
print(f"input video is NOT exist.{input_path}")
|
return False
|
if bg_video and not os.path.exists(bg_video):
|
print(f"input background video is NOT exist.{bg_video}")
|
return False
|
if bg_image and not os.path.exists(bg_image):
|
print(f"input background image is NOT exist.{bg_image}")
|
return False
|
|
try:
|
mask_ = generate_mask_from_first_frame(video_path=input_path,filter_type=filter_type,mask_type="file")
|
|
if mask_:
|
|
foreground_path, alpha_path = self.processor.process_video(
|
input_path = input_path,
|
mask_path = mask_,
|
output_path = output_path,
|
bg_video = bg_video,
|
bg_image = bg_image
|
)
|
return True
|
return False
|
|
except Exception as e:
|
traceback.extract_stack()
|
self.logger.error(f"[ error ], 发生错误: {e}")
|
return False
|
