waifuc/action/basic.py

from typing import Optional

from imgutils.data import load_image

from .base import ProcessAction
from ..model import ImageItem


class ModeConvertAction(ProcessAction):
    def __init__(self, mode='RGB', force_background: Optional[str] = 'white'):
        self.mode = mode
        self.force_background = force_background

    def process(self, item: ImageItem) -> ImageItem:
        image = load_image(item.image, mode=self.mode, force_background=self.force_background)
        return ImageItem(image, item.meta)