DataProcessingFramework

1
from enum import IntEnum
2

3

4
# types of resizing logic in Resizer
5
class ResizerModes(IntEnum):
6
    FIXED = 0
7
    MIN_SIZE = 1
8
    MAX_SIZE = 2
9

10

11
class Resizer:
12
    """
13
    Resizer class that used for different resizing strategies
14
    """
15

16
    def __init__(
17
        self,
18
        mode: ResizerModes,
19
        fixed_size: tuple[int, int] = (512, 512),
20
        size: int = 512,
21
        downscale_only: bool = True
22
    ):
23
        self.mode = mode
24
        assert self.mode in [ResizerModes.FIXED, ResizerModes.MIN_SIZE, ResizerModes.MAX_SIZE], \
25
            f"Invalid resizer mode: {self.mode}. Use mode from ResizerModes"
26

27
        self.fixed_size = fixed_size
28
        self.size = size
29
        self.downscale_only = downscale_only
30

31
    def get_new_size(self, width: int, height: int) -> tuple[int, int]:
32
        if self.mode == ResizerModes.FIXED:
33
            new_w, new_h = self.fixed_size
34
        elif self.mode == ResizerModes.MIN_SIZE:
35
            if self.downscale_only and min(height, width) < self.size:
36
                new_w, new_h = width, height
37
            elif height >= width:
38
                new_w = self.size
39
                new_h = int(self.size * height / width)
40
            else:
41
                new_h = self.size
42
                new_w = int(self.size * width / height)
43
        else:
44
            if self.downscale_only and max(height, width) < self.size:
45
                new_w, new_h = width, height
46
            elif width >= height:
47
                new_w = self.size
48
                new_h = int(self.size * height / width)
49
            else:
50
                new_h = self.size
51
                new_w = int(self.size * width / height)
52

53
        return new_w, new_h
54