from PIL import Image
import numpy as np
from preprocess.humanparsing.run_parsing import Parsing

parsing_model = Parsing(0)

class BackgroundProcessor:
    @classmethod
    def add_background(cls, human_img: Image, background_img: Image):
        
        human_img = human_img.convert("RGB")
        width = human_img.width
        height = human_img.height
        
        # Create mask image
        parsed_img, _ = parsing_model(human_img)
        mask_img = parsed_img.convert("L")
        mask_img = mask_img.resize((width, height))
        
        background_img = background_img.convert("RGB")
        background_img = background_img.resize((width, height))

        # Convert to numpy arrays
        human_np = np.array(human_img)
        mask_np = np.array(mask_img)
        background_np = np.array(background_img)

        # Ensure mask is 3-channel (RGB) for compatibility
        mask_np = np.stack((mask_np,) * 3, axis=-1)

        # Apply the mask to human_img
        human_with_background = np.where(mask_np > 0, human_np, background_np)

        # Convert back to PIL Image
        result_img = Image.fromarray(human_with_background.astype('uint8'))

        # Return or save the result
        return result_img