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