waifuc/action/augument.py

import os.path
import random
from typing import Iterator, Optional, Tuple

from PIL import ImageOps
from hbutils.random import random_sha1

from .base import BaseAction
from ..model import ImageItem


class BaseRandomAction(BaseAction):
    def __init__(self, seed=None):
        self.seed = seed
        self.random = random.Random(self.seed)

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        raise NotImplementedError  # pragma: no cover

    def reset(self):
        self.random = random.Random(self.seed)


class RandomChoiceAction(BaseRandomAction):
    def __init__(self, p=0.5, seed=None):
        BaseRandomAction.__init__(self, seed)
        self.p = p

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        if self.random.random() <= self.p:
            yield item


class RandomFilenameAction(BaseRandomAction):
    def __init__(self, ext: Optional[str] = '.png', seed=None):
        BaseRandomAction.__init__(self, seed)
        self.ext = ext

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        if 'filename' in item.meta:
            ext = self.ext or os.path.splitext(os.path.basename(item.meta['filename']))[0]
        else:
            if self.ext:
                ext = self.ext
            else:
                raise NameError(f'Extension (ext) must be specified '
                                f'when filename not in metadata of image item - {item!r}.')

        filename = random_sha1(rnd=self.random) + ext
        yield ImageItem(item.image, {**item.meta, 'filename': filename})


class MirrorAction(BaseAction):
    def __init__(self, names: Tuple[str, str] = ('origin', 'mirror')):
        self.origin_name, self.mirror_name = names

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        if 'filename' in item.meta:
            filebody, ext = os.path.splitext(item.meta['filename'])
            yield ImageItem(item.image, {**item.meta, 'filename': f'{filebody}_{self.origin_name}{ext}'})
            yield ImageItem(ImageOps.mirror(item.image),
                            {**item.meta, 'filename': f'{filebody}_{self.mirror_name}{ext}'})
        else:
            yield ImageItem(item.image, item.meta)
            yield ImageItem(ImageOps.mirror(item.image), item.meta)

    def reset(self):
        pass