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 vkit.engine.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(self, run_config: ImageEngineRunConfig, rng: RandomGenerator) -> Image:
 66        image_file = rng_choice(rng, self.image_files)
 67        image = Image.from_file(image_file)
 68
 69        if self.init_config.target_image_mode:
 70            image = image.to_target_mode_image(self.init_config.target_image_mode)
 71
 72        if run_config.disable_resizing:
 73            assert run_config.height == 0 and run_config.width == 0
 74            return image
 75
 76        height = run_config.height
 77        width = run_config.width
 78        if not self.init_config.force_resize and height <= image.height and width <= image.width:
 79            # Select a part of image.
 80            up = rng.integers(0, image.height - height + 1)
 81            left = rng.integers(0, image.width - width + 1)
 82            box = Box(
 83                up=up,
 84                down=up + height - 1,
 85                left=left,
 86                right=left + width - 1,
 87            )
 88            image = box.extract_image(image)
 89
 90        else:
 91            # Resize image.
 92            image = image.to_resized_image(
 93                resized_height=height,
 94                resized_width=width,
 95            )
 96
 97        return image
 98
 99
100image_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(self, run_config: ImageEngineRunConfig, rng: RandomGenerator) -> Image:
67        image_file = rng_choice(rng, self.image_files)
68        image = Image.from_file(image_file)
69
70        if self.init_config.target_image_mode:
71            image = image.to_target_mode_image(self.init_config.target_image_mode)
72
73        if run_config.disable_resizing:
74            assert run_config.height == 0 and run_config.width == 0
75            return image
76
77        height = run_config.height
78        width = run_config.width
79        if not self.init_config.force_resize and height <= image.height and width <= image.width:
80            # Select a part of image.
81            up = rng.integers(0, image.height - height + 1)
82            left = rng.integers(0, image.width - width + 1)
83            box = Box(
84                up=up,
85                down=up + height - 1,
86                left=left,
87                right=left + width - 1,
88            )
89            image = box.extract_image(image)
90
91        else:
92            # Resize image.
93            image = image.to_resized_image(
94                resized_height=height,
95                resized_width=width,
96            )
97
98        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: numpy.random._generator.Generator) -> vkit.element.image.Image: View Source

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