Spaces:
Sleeping
Sleeping
| from cil.optimisation.utilities import AlgorithmDiagnostics | |
| import numpy as np | |
| from bm3d import bm3d, BM3DStages | |
| from cil.optimisation.functions import Function | |
| class StoppingCriterionTime(AlgorithmDiagnostics): | |
| def __init__(self, time): | |
| self.time = time | |
| super(StoppingCriterionTime, self).__init__(verbose=0) | |
| self.should_stop = False | |
| def _should_stop(self): | |
| return self.should_stop | |
| def __call__(self, algo): | |
| if algo.iteration==0: | |
| algo.should_stop = self._should_stop | |
| stop_crit = np.sum(algo.timing)>self.time | |
| if stop_crit: | |
| self.should_stop = True | |
| print("Stop at {} time {}".format(algo.iteration, np.sum(algo.timing))) | |
| class BM3DFunction(Function): | |
| """ | |
| PnP 'regulariser' whose proximal applies BM3D denoising. | |
| In PnP-ISTA/FISTA we typically use a FIXED BM3D sigma (regularization strength), | |
| independent of the gradient step-size tau. | |
| Optionally apply damping: (1-gamma) z + gamma * BM3D(z). | |
| """ | |
| def __init__(self, sigma, gamma=1.0, profile="np", | |
| stage_arg=BM3DStages.ALL_STAGES, positivity=True): | |
| self.sigma = float(sigma) # BM3D noise parameter | |
| self.gamma = float(gamma) # damping in (0,1] | |
| if not (0.0 < self.gamma <= 1.0): | |
| raise ValueError("gamma must be in (0,1].") | |
| self.profile = profile | |
| self.stage_arg = stage_arg | |
| self.positivity = positivity | |
| super().__init__() | |
| def __call__(self, x): | |
| return 0.0 | |
| def convex_conjugate(self, x): | |
| return 0.0 | |
| def _denoise(self, znp: np.ndarray) -> np.ndarray: | |
| z = np.asarray(znp, dtype=np.float32) | |
| # BM3D expects sigma as noise std (same units as the image) | |
| return bm3d(z, sigma_psd=self.sigma, | |
| profile=self.profile, | |
| stage_arg=self.stage_arg).astype(np.float32) | |
| def proximal(self, x, tau, out=None): | |
| z = x.array.astype(np.float32, copy=False) | |
| d = self._denoise(z) | |
| # damping/relaxation (recommended if you see oscillations) | |
| u = (1.0 - self.gamma) * z + self.gamma * d | |
| if self.positivity: | |
| u = np.maximum(u, 0.0) | |
| if out is None: | |
| out = x * 0.0 | |
| out.fill(u) | |
| return out | |
| def create_circular_mask(h, w, center=None, radius=None): | |
| if center is None: | |
| center = (int(w/2), int(h/2)) | |
| if radius is None: | |
| radius = min(center[0], center[1], w-center[0], h-center[1]) | |
| Y, X = np.ogrid[:h, :w] | |
| dist_from_center = np.sqrt((X - center[0])**2 + (Y-center[1])**2) | |
| mask = dist_from_center <= radius | |
| return mask | |
| class StoppingCriterion(AlgorithmDiagnostics): | |
| def __init__(self, epsilon, epochs=None): | |
| self.epsilon = epsilon | |
| self.epochs = epochs | |
| super().__init__(verbose=0) | |
| self.should_stop = False | |
| self.rse_reached = False | |
| def _should_stop(self): | |
| return self.should_stop | |
| def __call__(self, algo): | |
| if algo.iteration == 0: | |
| algo.should_stop = self._should_stop | |
| stop_rse = (algo.rse[-1] <= self.epsilon) | |
| stop_epochs = False | |
| if self.epochs is not None: | |
| try: | |
| dp = algo.f.data_passes | |
| dp_last = dp[-1] if hasattr(dp, "__len__") else dp | |
| stop_epochs = (dp_last > self.epochs) | |
| except AttributeError: | |
| stop_epochs = False | |
| stop = stop_rse or stop_epochs | |
| if algo.iteration < algo.max_iteration: | |
| if stop: | |
| self.rse_reached = stop_rse | |
| self.should_stop = True | |
| print(f"Accuracy reached at {algo.iteration}, time = {np.sum(algo.timing):.4f}, NRSE = {algo.rse[-1]:.4e}") | |
| else: | |
| print(f"Failed to reach accuracy. Stop at {algo.iteration}, time = {np.sum(algo.timing):.4f}, NRSE = {algo.rse[-1]:.4e}") | |