88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
import numpy as np
|
|
from torchvision import transforms
|
|
from PIL import Image
|
|
|
|
|
|
def transform_image(image: Image) -> Image:
|
|
im_shape = 64
|
|
resize_transforms = transforms.Compose([
|
|
transforms.Resize(size=im_shape),
|
|
transforms.CenterCrop(size=(im_shape, im_shape))
|
|
])
|
|
return resize_transforms(image)
|
|
|
|
|
|
def to_grayscale(pil_image: np.ndarray) -> np.ndarray:
|
|
if pil_image.ndim == 2:
|
|
return pil_image.copy()[None]
|
|
if pil_image.ndim != 3:
|
|
raise ValueError("image must have either shape (H, W) or (H, W, 3)")
|
|
if pil_image.shape[2] != 3:
|
|
raise ValueError(f"image has shape (H, W, {pil_image.shape[2]}), but it should have (H, W, 3)")
|
|
|
|
rgb = pil_image / 255
|
|
rgb_linear = np.where(
|
|
rgb < 0.04045,
|
|
rgb / 12.92,
|
|
((rgb + 0.055) / 1.055) ** 2.4
|
|
)
|
|
grayscale_linear = 0.2126 * rgb_linear[..., 0] + 0.7152 * rgb_linear[..., 1] + 0.0722 * rgb_linear[..., 2]
|
|
|
|
grayscale = np.where(
|
|
grayscale_linear < 0.0031308,
|
|
12.92 * grayscale_linear,
|
|
1.055 * grayscale_linear ** (1 / 2.4) - 0.055
|
|
)
|
|
grayscale = grayscale * 255
|
|
|
|
if np.issubdtype(pil_image.dtype, np.integer):
|
|
grayscale = np.round(grayscale)
|
|
return grayscale.astype(pil_image.dtype)[None]
|
|
|
|
|
|
def prepare_image(image: np.ndarray, x: int, y: int, width: int, height: int, size: int) -> \
|
|
tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
if image.ndim < 3 or image.shape[-3] != 1:
|
|
# This is actually more general than the assignment specification
|
|
raise ValueError("image must have shape (..., 1, H, W)")
|
|
if width < 2 or height < 2 or size < 2:
|
|
raise ValueError("width/height/size must be >= 2")
|
|
if x < 0 or (x + width) > image.shape[-1]:
|
|
raise ValueError(f"x={x} and width={width} do not fit into the image width={image.shape[-1]}")
|
|
if y < 0 or (y + height) > image.shape[-2]:
|
|
raise ValueError(f"y={y} and height={height} do not fit into the image height={image.shape[-2]}")
|
|
|
|
# The (height, width) slices to extract the area that should be pixelated. Since we
|
|
# need this multiple times, specify the slices explicitly instead of using [:] notation
|
|
area = (..., slice(y, y + height), slice(x, x + width))
|
|
|
|
# This returns already a copy, so we are independent of "image"
|
|
pixelated_image = pixelate(image, x, y, width, height, size)
|
|
|
|
known_array = np.full_like(image, fill_value=False, dtype=bool)
|
|
known_array[area] = True
|
|
|
|
# Create a copy to avoid that "target_array" and "image" point to the same array
|
|
# target_array = image[area].copy()
|
|
stacked_target_array = np.full((1, 64, 64), 0)
|
|
stacked_target_array[area] = image[area].copy()
|
|
target_array = stacked_target_array
|
|
|
|
return pixelated_image, known_array, target_array
|
|
|
|
|
|
def pixelate(image: np.ndarray, x: int, y: int, width: int, height: int, size: int) -> np.ndarray:
|
|
# Need a copy since we overwrite data directly
|
|
image = image.copy()
|
|
curr_x = x
|
|
|
|
while curr_x < x + width:
|
|
curr_y = y
|
|
while curr_y < y + height:
|
|
block = (..., slice(curr_y, min(curr_y + size, y + height)), slice(curr_x, min(curr_x + size, x + width)))
|
|
image[block] = image[block].mean()
|
|
curr_y += size
|
|
curr_x += size
|
|
|
|
return image
|