134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
from logging import getLogger
|
|
from typing import Optional, Tuple
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from ..params import ImageParams, SizeChart, StageParams
|
|
from ..server import ServerContext
|
|
from ..worker import ProgressCallback, WorkerContext
|
|
from .base import BaseStage
|
|
from .result import StageResult
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
class BlendDenoiseLocalStdStage(BaseStage):
|
|
"""
|
|
Experimental stage to blend and denoise images using local means compared to local standard deviation.
|
|
Very slow.
|
|
"""
|
|
|
|
max_tile = SizeChart.max
|
|
|
|
def run(
|
|
self,
|
|
_worker: WorkerContext,
|
|
_server: ServerContext,
|
|
_stage: StageParams,
|
|
_params: ImageParams,
|
|
sources: StageResult,
|
|
*,
|
|
strength: int = 5,
|
|
range: int = 4,
|
|
stage_source: Optional[Image.Image] = None,
|
|
callback: Optional[ProgressCallback] = None,
|
|
**kwargs,
|
|
) -> StageResult:
|
|
logger.info("denoising source images")
|
|
|
|
return StageResult.from_arrays(
|
|
[
|
|
remove_noise(source, threshold=strength, deviation=range)[0]
|
|
for source in sources.as_arrays()
|
|
],
|
|
metadata=sources.metadata,
|
|
)
|
|
|
|
|
|
def downscale_image(image: np.ndarray, scale: int = 2):
|
|
result_image = np.zeros(
|
|
(image.shape[0] // scale, image.shape[1] // scale), dtype=np.uint8
|
|
)
|
|
|
|
for i in range(0, image.shape[0] - scale, scale):
|
|
for j in range(0, image.shape[1] - scale, scale):
|
|
# Average the four neighboring pixels
|
|
pixel_average = np.mean(image[i : i + scale, j : j + scale], axis=(0, 1))
|
|
result_image[i // scale, j // scale] = pixel_average.astype(np.uint8)
|
|
|
|
return result_image
|
|
|
|
|
|
def replace_noise(region: np.ndarray, threshold: int, deviation: float, op=np.median):
|
|
# Identify stray pixels (brightness significantly deviates from surrounding pixels)
|
|
central_pixel = np.mean(region[2:4, 2:4])
|
|
|
|
region_normal = op(region)
|
|
region_deviation = np.std(region)
|
|
diff = np.abs(central_pixel - region_normal)
|
|
|
|
# If the whole region is fairly consistent but the central pixel deviates significantly,
|
|
if diff > (region_deviation + threshold) and diff < (
|
|
region_deviation + threshold * deviation
|
|
):
|
|
surrounding_pixels = region[region != central_pixel]
|
|
surrounding_median = op(surrounding_pixels)
|
|
# replace it with the median of surrounding pixels
|
|
region[1, 1] = surrounding_median
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def remove_noise(
|
|
image: np.ndarray,
|
|
threshold: int,
|
|
deviation: float,
|
|
region_size: Tuple[int, int] = (6, 6),
|
|
):
|
|
# Create a copy of the original image to store the result
|
|
result_image = np.copy(image)
|
|
result_mask = np.zeros_like(image)
|
|
|
|
# Iterate over regions in each channel
|
|
i_inc = region_size[0] // 2
|
|
j_inc = region_size[1] // 2
|
|
|
|
for i in range(i_inc, image.shape[0] - i_inc, 1):
|
|
for j in range(j_inc, image.shape[1] - j_inc, 1):
|
|
i_min = i - (region_size[0] // 2)
|
|
i_max = i + (region_size[0] // 2)
|
|
j_min = j - (region_size[1] // 2)
|
|
j_max = j + (region_size[1] // 2)
|
|
# print(i_min, i_max, j_min, j_max)
|
|
|
|
# skip if the central pixels have already been masked by a previous artifact
|
|
if np.any(result_mask[i - 1 : i + 1, j - 1 : j + 1] > 0):
|
|
continue
|
|
|
|
# Extract region from each channel
|
|
region_red = image[i_min:i_max, j_min:j_max, 0]
|
|
region_green = image[i_min:i_max, j_min:j_max, 1]
|
|
region_blue = image[i_min:i_max, j_min:j_max, 2]
|
|
|
|
replaced = any(
|
|
[
|
|
replace_noise(region_red, threshold, deviation),
|
|
replace_noise(region_green, threshold, deviation),
|
|
]
|
|
)
|
|
|
|
# apply the noise removal function to each channel
|
|
if replaced:
|
|
# assign the processed region back to the result image
|
|
result_image[i - 1 : i + 1, j - 1 : j + 1, 0] = region_red[1, 1]
|
|
result_image[i - 1 : i + 1, j - 1 : j + 1, 1] = region_green[1, 1]
|
|
result_image[i - 1 : i + 1, j - 1 : j + 1, 2] = region_blue[1, 1]
|
|
|
|
result_mask[i - 1 : i + 1, j - 1 : j + 1, 0] = 1
|
|
result_mask[i - 1 : i + 1, j - 1 : j + 1, 1] = 1
|
|
result_mask[i - 1 : i + 1, j - 1 : j + 1, 2] = 1
|
|
|
|
return (result_image, result_mask * 255)
|