Spaces:
Runtime error
Runtime error
File size: 2,301 Bytes
4f8ad24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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
|