utils.py

# -*- coding: utf-8 -*-
# @Author: 
# @Date: 2025-06-11 19:01:25
# @LastEditTime: 2025-06-13 20:43:05
# @FilePath: /image_denoise/utils.py
# @Description: 通用的一些算子

import numpy as np
import torch
import torch.nn as nn

def gradient_x(x):
	"""计算x方向梯度"""
	grad = np.diff(x, axis=1)
	# 边界处理：循环边界条件
	boundary = x[:, 0:1] - x[:, -1:]
	return np.concatenate([grad, boundary], axis=1)

def gradient_y(y):
	"""计算y方向梯度"""
	grad = np.diff(y, axis=0)
	# 边界处理：循环边界条件
	boundary = y[0:1, :] - y[-1:, :]
	return np.concatenate([grad, boundary], axis=0)

def divergence(vec_x, vec_y):
	"""计算散度 (梯度的转置算子)"""
	# x方向散度
	div_x_inner = -np.diff(vec_x, axis=1)
	div_x_boundary = vec_x[:, -1:] - vec_x[:, 0:1]
	div_x = np.concatenate([div_x_boundary, div_x_inner], axis=1)
	
	# y方向散度
	div_y_inner = -np.diff(vec_y, axis=0)
	div_y_boundary = vec_y[-1:, :] - vec_y[0:1, :]
	div_y = np.concatenate([div_y_boundary, div_y_inner], axis=0)
	
	return div_x + div_y

def zero_pad(image, shape, position='corner'):
	"""
	Extends image to a certain size with zeros

	Parameters
	----------
	image: real 2d `numpy.ndarray`
		Input image
	shape: tuple of int
		Desired output shape of the image
	position : str, optional
		The position of the input image in the output one:
			* 'corner'
				top-left corner (default)
			* 'center'
				centered

	Returns
	-------
	padded_img: real `numpy.ndarray`
		The zero-padded image

	"""
	shape = np.asarray(shape, dtype=int)
	imshape = np.asarray(image.shape, dtype=int)

	if np.all(imshape == shape):
		return image

	if np.any(shape <= 0):
		raise ValueError("ZERO_PAD: null or negative shape given")

	dshape = shape - imshape
	if np.any(dshape < 0):
		raise ValueError("ZERO_PAD: target size smaller than source one")

	pad_img = np.zeros(shape, dtype=image.dtype)

	idx, idy = np.indices(imshape)

	if position == 'center':
		if np.any(dshape % 2 != 0):
			raise ValueError("ZERO_PAD: source and target shapes have different parity.")
		offx, offy = dshape // 2
	else:
		offx, offy = (0, 0)

	pad_img[idx + offx, idy + offy] = image

	return pad_img

def psf2otf(psf, shape):
	"""
	Convert point-spread function to optical transfer function.
	------------------------------------------------------------
	Compute the Fast Fourier Transform (FFT) of the point-spread
	function (PSF) array and creates the optical transfer function (OTF)
	array that is not influenced by the PSF off-centering.
	By default, the OTF array is the same size as the PSF array.

	To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
	post-pads the PSF array (down or to the right) with zeros to match
	dimensions specified in OUTSIZE, then circularly shifts the values of
	the PSF array up (or to the left) until the central pixel reaches (1,1)
	position.

	Parameters
	----------
	psf : `numpy.ndarray`
		PSF array
	shape : int
		Output shape of the OTF array

	Returns
	-------
	otf : `numpy.ndarray`
		OTF array

	Notes
	-----
	Adapted from MATLAB psf2otf function
	Arrays of higher dimension that 2D are also supported
	"""
	if np.all(psf == 0):
		return np.zeros_like(psf)

	inshape = psf.shape
	# Pad the PSF to outsize
	psf = zero_pad(psf, shape, position='corner')

	# Circularly shift OTF so that the 'center' of the PSF is
	# [0,0] element of the array
	for axis, axis_size in enumerate(inshape):
		psf = np.roll(psf, -int(axis_size / 2), axis=axis)

	# Compute the OTF
	otf = np.fft.fftn(psf)

	# Estimate the rough number of operations involved in the FFT
	# and discard the PSF imaginary part if within roundoff error
	# roundoff error  = machine epsilon = sys.float_info.epsilon
	# or np.finfo().eps
	n_ops = np.sum(psf.size * np.log2(psf.shape))
	otf = np.real_if_close(otf, tol=n_ops)

	return otf

def soft_threshold(x: np.ndarray, tau: float) -> np.ndarray:
	"""
	应用软阈值操作

	Args:
		x (np.ndarray): 输入数组
		tau (float): 阈值

	Returns:
		np.ndarray: 软阈值处理后的数组
	"""
	return np.sign(x) * np.maximum(np.abs(x) - tau, 0)

# 测试soft_threshold函数
if __name__ == "__main__":
		x = np.array([[1, -2, 3], [-4, 5, -6]])
		tau = 2.0
		result = soft_threshold(x, tau)
		print("Soft Threshold Result:\n", result)