36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from random import randint
|
|
from .point import Point
|
|
|
|
|
|
def tensor_to_pil(image_tensor: torch.tensor) -> Image:
|
|
image_np = image_tensor.cpu().numpy()
|
|
image_pil = Image.fromarray((image_np.squeeze(0) * 255).astype(np.uint8))
|
|
return image_pil.convert("RGBA")
|
|
|
|
|
|
def pil_to_tensor(image_pil: Image) -> torch.Tensor:
|
|
image_tensor_out = torch.tensor(
|
|
np.array(image_pil).astype(np.float32) / 255.0)
|
|
return torch.unsqueeze(image_tensor_out, 0)
|
|
|
|
|
|
def generate_noise(image, point1: Point, point2: Point, padding=0):
|
|
def noise_color(i: int, j: int, color_index: int) -> int:
|
|
noise = randint(-64, 64)
|
|
return max(0, min(pixels[i, j][color_index] + noise, 255))
|
|
|
|
pixels = image.load()
|
|
for i in range(point1.x - padding, point2.x + padding):
|
|
for j in range(point1.y - padding, point2.y + padding):
|
|
try:
|
|
pixels[i, j] = (
|
|
noise_color(i, j, 0),
|
|
noise_color(i, j, 1),
|
|
noise_color(i, j, 2)
|
|
)
|
|
except IndexError:
|
|
pass
|