Spaces:
Runtime error
Runtime error
File size: 2,341 Bytes
c19ca42 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
from enum import Enum
from typing import Tuple
import numpy as np
from PIL import Image
from ..utils.utils import get_h_w_c
from .image_utils import FillColor, convert_to_BGRA, normalize, to_uint8
class InterpolationMethod(Enum):
AUTO = -1
NEAREST = 0
LANCZOS = 1
LINEAR = 2
CUBIC = 3
BOX = 4
class RotationInterpolationMethod(Enum):
CUBIC = InterpolationMethod.CUBIC.value
LINEAR = InterpolationMethod.LINEAR.value
NEAREST = InterpolationMethod.NEAREST.value
@property
def interpolation_method(self) -> InterpolationMethod:
return InterpolationMethod(self.value)
INTERPOLATION_METHODS_MAP = {
InterpolationMethod.NEAREST: Image.NEAREST,
InterpolationMethod.BOX: Image.BOX,
InterpolationMethod.LINEAR: Image.BILINEAR,
InterpolationMethod.CUBIC: Image.BICUBIC,
InterpolationMethod.LANCZOS: Image.LANCZOS,
}
class RotateSizeChange(Enum):
EXPAND = 1
CROP = 0
def resize(
img: np.ndarray, out_dims: Tuple[int, int], interpolation: InterpolationMethod
) -> np.ndarray:
"""Perform PIL resize"""
if interpolation == InterpolationMethod.AUTO:
# automatically chose a method that works
new_w, new_h = out_dims
old_h, old_w, _ = get_h_w_c(img)
if new_w > old_w or new_h > old_h:
interpolation = InterpolationMethod.LANCZOS
else:
interpolation = InterpolationMethod.BOX
resample = INTERPOLATION_METHODS_MAP[interpolation]
pimg = Image.fromarray(to_uint8(img, normalized=True))
pimg = pimg.resize(out_dims, resample=resample) # type: ignore
return normalize(np.array(pimg))
def rotate(
img: np.ndarray,
angle: float,
interpolation: RotationInterpolationMethod,
expand: RotateSizeChange,
fill: FillColor,
) -> np.ndarray:
"""Perform PIL rotate"""
c = get_h_w_c(img)[2]
if fill == FillColor.TRANSPARENT:
img = convert_to_BGRA(img, c)
fill_color = tuple([x * 255 for x in fill.get_color(c)])
resample = INTERPOLATION_METHODS_MAP[interpolation.interpolation_method]
pimg = Image.fromarray(to_uint8(img, normalized=True))
pimg = pimg.rotate(
angle,
resample=resample, # type: ignore
expand=bool(expand.value),
fillcolor=fill_color,
)
return normalize(np.array(pimg))
|