vkit.engine.image.selector

View Source

  1# Copyright 2022 vkit-x Administrator. All Rights Reserved.
  2#
  3# This project (vkit-x/vkit) is dual-licensed under commercial and SSPL licenses.
  4#
  5# The commercial license gives you the full rights to create and distribute software
  6# on your own terms without any SSPL license obligations. For more information,
  7# please see the "LICENSE_COMMERCIAL.txt" file.
  8#
  9# This project is also available under Server Side Public License (SSPL).
 10# The SSPL licensing is ideal for use cases such as open source projects with
 11# SSPL distribution, student/academic purposes, hobby projects, internal research
 12# projects without external distribution, or other projects where all SSPL
 13# obligations can be met. For more information, please see the "LICENSE_SSPL.txt" file.
 14from typing import List, Optional, Sequence
 15from os import PathLike
 16
 17import attrs
 18from numpy.random import Generator as RandomGenerator
 19import iolite as io
 20
 21from vkit.utility import rng_choice
 22from vkit.element import Image, ImageMode, Box
 23from ..interface import (
 24    Engine,
 25    EngineExecutorFactory,
 26    NoneTypeEngineInitResource,
 27)
 28from .type import ImageEngineRunConfig
 29
 30
 31@attrs.define
 32class ImageSelectorEngineInitConfig:
 33    image_folders: Sequence[str]
 34    target_image_mode: Optional[ImageMode] = ImageMode.RGB
 35    force_resize: bool = False
 36
 37
 38class ImageSelectorEngine(
 39    Engine[
 40        ImageSelectorEngineInitConfig,
 41        NoneTypeEngineInitResource,
 42        ImageEngineRunConfig,
 43        Image,
 44    ]
 45):  # yapf: disable
 46
 47    @classmethod
 48    def get_type_name(cls) -> str:
 49        return 'selector'
 50
 51    def __init__(
 52        self,
 53        init_config: ImageSelectorEngineInitConfig,
 54        init_resource: Optional[NoneTypeEngineInitResource] = None,
 55    ):
 56        super().__init__(init_config, init_resource)
 57
 58        self.image_files: List[PathLike] = []
 59        for image_folder in self.init_config.image_folders:
 60            image_fd = io.folder(image_folder, expandvars=True, exists=True)
 61            for ext in ['jpg', 'jpeg', 'png']:
 62                for new_ext in [ext, ext.upper()]:
 63                    self.image_files.extend(image_fd.glob(f'**/*.{new_ext}'))
 64
 65    def run(
 66        self,
 67        run_config: ImageEngineRunConfig,
 68        rng: Optional[RandomGenerator] = None,
 69    ) -> Image:
 70        assert rng is not None
 71
 72        image_file = rng_choice(rng, self.image_files)
 73        image = Image.from_file(image_file)
 74
 75        if self.init_config.target_image_mode:
 76            image = image.to_target_mode_image(self.init_config.target_image_mode)
 77
 78        if run_config.disable_resizing:
 79            assert run_config.height == 0 and run_config.width == 0
 80            return image
 81
 82        height = run_config.height
 83        width = run_config.width
 84        if not self.init_config.force_resize and height <= image.height and width <= image.width:
 85            # Select a part of image.
 86            up = rng.integers(0, image.height - height + 1)
 87            left = rng.integers(0, image.width - width + 1)
 88            box = Box(
 89                up=up,
 90                down=up + height - 1,
 91                left=left,
 92                right=left + width - 1,
 93            )
 94            image = box.extract_image(image)
 95
 96        else:
 97            # Resize image.
 98            image = image.to_resized_image(
 99                resized_height=height,
100                resized_width=width,
101            )
102
103        return image
104
105
106image_selector_engine_executor_factory = EngineExecutorFactory(ImageSelectorEngine)

class ImageSelectorEngineInitConfig: View Source

33class ImageSelectorEngineInitConfig:
34    image_folders: Sequence[str]
35    target_image_mode: Optional[ImageMode] = ImageMode.RGB
36    force_resize: bool = False

ImageSelectorEngineInitConfig( image_folders: Sequence[str], target_image_mode: Union[vkit.element.image.ImageMode, NoneType] = <ImageMode.RGB: 'rgb'>, force_resize: bool = False) View Source

2def __init__(self, image_folders, target_image_mode=attr_dict['target_image_mode'].default, force_resize=attr_dict['force_resize'].default):
3    self.image_folders = image_folders
4    self.target_image_mode = target_image_mode
5    self.force_resize = force_resize

Method generated by attrs for class ImageSelectorEngineInitConfig.

class ImageSelectorEngine(vkit.engine.interface.Engine[vkit.engine.image.selector.ImageSelectorEngineInitConfig, vkit.engine.interface.NoneTypeEngineInitResource, vkit.engine.image.type.ImageEngineRunConfig, vkit.element.image.Image]): View Source

 39class ImageSelectorEngine(
 40    Engine[
 41        ImageSelectorEngineInitConfig,
 42        NoneTypeEngineInitResource,
 43        ImageEngineRunConfig,
 44        Image,
 45    ]
 46):  # yapf: disable
 47
 48    @classmethod
 49    def get_type_name(cls) -> str:
 50        return 'selector'
 51
 52    def __init__(
 53        self,
 54        init_config: ImageSelectorEngineInitConfig,
 55        init_resource: Optional[NoneTypeEngineInitResource] = None,
 56    ):
 57        super().__init__(init_config, init_resource)
 58
 59        self.image_files: List[PathLike] = []
 60        for image_folder in self.init_config.image_folders:
 61            image_fd = io.folder(image_folder, expandvars=True, exists=True)
 62            for ext in ['jpg', 'jpeg', 'png']:
 63                for new_ext in [ext, ext.upper()]:
 64                    self.image_files.extend(image_fd.glob(f'**/*.{new_ext}'))
 65
 66    def run(
 67        self,
 68        run_config: ImageEngineRunConfig,
 69        rng: Optional[RandomGenerator] = None,
 70    ) -> Image:
 71        assert rng is not None
 72
 73        image_file = rng_choice(rng, self.image_files)
 74        image = Image.from_file(image_file)
 75
 76        if self.init_config.target_image_mode:
 77            image = image.to_target_mode_image(self.init_config.target_image_mode)
 78
 79        if run_config.disable_resizing:
 80            assert run_config.height == 0 and run_config.width == 0
 81            return image
 82
 83        height = run_config.height
 84        width = run_config.width
 85        if not self.init_config.force_resize and height <= image.height and width <= image.width:
 86            # Select a part of image.
 87            up = rng.integers(0, image.height - height + 1)
 88            left = rng.integers(0, image.width - width + 1)
 89            box = Box(
 90                up=up,
 91                down=up + height - 1,
 92                left=left,
 93                right=left + width - 1,
 94            )
 95            image = box.extract_image(image)
 96
 97        else:
 98            # Resize image.
 99            image = image.to_resized_image(
100                resized_height=height,
101                resized_width=width,
102            )
103
104        return image

Abstract base class for generic types.

A generic type is typically declared by inheriting from this class parameterized with one or more type variables. For example, a generic mapping type might be defined as::

class Mapping(Generic[KT, VT]): def __getitem__(self, key: KT) -> VT: ... # Etc.

This class can then be used as follows::

def lookup_name(mapping: Mapping[KT, VT], key: KT, default: VT) -> VT: try: return mapping[key] except KeyError: return default

ImageSelectorEngine( init_config: vkit.engine.image.selector.ImageSelectorEngineInitConfig, init_resource: Union[vkit.engine.interface.NoneTypeEngineInitResource, NoneType] = None) View Source

52    def __init__(
53        self,
54        init_config: ImageSelectorEngineInitConfig,
55        init_resource: Optional[NoneTypeEngineInitResource] = None,
56    ):
57        super().__init__(init_config, init_resource)
58
59        self.image_files: List[PathLike] = []
60        for image_folder in self.init_config.image_folders:
61            image_fd = io.folder(image_folder, expandvars=True, exists=True)
62            for ext in ['jpg', 'jpeg', 'png']:
63                for new_ext in [ext, ext.upper()]:
64                    self.image_files.extend(image_fd.glob(f'**/*.{new_ext}'))

@classmethod

def get_type_name(cls) -> str: View Source

48    @classmethod
49    def get_type_name(cls) -> str:
50        return 'selector'

def run( self, run_config: vkit.engine.image.type.ImageEngineRunConfig, rng: Union[numpy.random._generator.Generator, NoneType] = None) -> vkit.element.image.Image: View Source

 66    def run(
 67        self,
 68        run_config: ImageEngineRunConfig,
 69        rng: Optional[RandomGenerator] = None,
 70    ) -> Image:
 71        assert rng is not None
 72
 73        image_file = rng_choice(rng, self.image_files)
 74        image = Image.from_file(image_file)
 75
 76        if self.init_config.target_image_mode:
 77            image = image.to_target_mode_image(self.init_config.target_image_mode)
 78
 79        if run_config.disable_resizing:
 80            assert run_config.height == 0 and run_config.width == 0
 81            return image
 82
 83        height = run_config.height
 84        width = run_config.width
 85        if not self.init_config.force_resize and height <= image.height and width <= image.width:
 86            # Select a part of image.
 87            up = rng.integers(0, image.height - height + 1)
 88            left = rng.integers(0, image.width - width + 1)
 89            box = Box(
 90                up=up,
 91                down=up + height - 1,
 92                left=left,
 93                right=left + width - 1,
 94            )
 95            image = box.extract_image(image)
 96
 97        else:
 98            # Resize image.
 99            image = image.to_resized_image(
100                resized_height=height,
101                resized_width=width,
102            )
103
104        return image