diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 5f0d64b..0797ddb 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -1,2 +1,4 @@ * Ask questions about the prompt to clarify any uncertainties. * Ensure that new code compiles without warnings or errors. +* Prioritize making new code efficient and compact +* Prefer file scoped namespaces \ No newline at end of file diff --git a/.gitignore b/.gitignore index 23b2679..5b5e627 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,5 @@ PhotoLocatorSetup/*.pfx PhotoLocatorSetup/AppPackages/ PhotoLocatorSetup/BundleArtifacts/ TestResults/ +__pycache__/ publish.bat diff --git a/.vscode/launch.json b/.vscode/launch.json index 42c2f7e..5ceaaa5 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -6,9 +6,9 @@ "name": "C#: PhotoLocator Debug", "type": "coreclr", "request": "launch", - "program": "${workspaceFolder}/PhotoLocator/bin/Debug/net8.0-windows/PhotoLocator.exe", + "program": "${workspaceFolder}/PhotoLocator/bin/Debug/net10.0-windows/PhotoLocator.exe", "args": [], - "cwd": "${workspaceFolder}/PhotoLocator/bin/Debug/net8.0-windows/", + "cwd": "${workspaceFolder}/PhotoLocator/bin/Debug/net10.0-windows/", "stopAtEntry": false, "preLaunchTask": "build" } diff --git a/PhotoLocator/AI/create_minimal_denoise_test_model.py b/PhotoLocator/AI/create_minimal_denoise_test_model.py new file mode 100644 index 0000000..7ad8abb --- /dev/null +++ b/PhotoLocator/AI/create_minimal_denoise_test_model.py @@ -0,0 +1,41 @@ +from pathlib import Path + +import onnx +from onnx import TensorProto, helper, numpy_helper + + +def create_identity_model(output_path: Path): + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", "channels", "height", "width"]) + output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", "channels", "height", "width"]) + + identity_node = helper.make_node( + "Identity", + inputs=["input"], + outputs=["output"], + name="identity_test", + ) + + graph = helper.make_graph( + [identity_node], + "minimal_denoise_test_graph", + [input_tensor], + [output_tensor], + ) + + model = helper.make_model(graph, producer_name="PhotoLocatorTest") + model.opset_import[0].version = 17 + model.ir_version = 13 + + output_path.parent.mkdir(parents=True, exist_ok=True) + onnx.save(model, str(output_path)) + print(f"Saved minimal ONNX smoke-test model to {output_path}") + + +def main(): + repo_root = Path(__file__).resolve().parents[2] + output_path = repo_root / "PhotoLocatorTest" / "TestData" / "dncnn_test.onnx" + create_identity_model(output_path) + + +if __name__ == "__main__": + main() diff --git a/PhotoLocator/AI/denoise_dncnn.py b/PhotoLocator/AI/denoise_dncnn.py new file mode 100644 index 0000000..511e380 --- /dev/null +++ b/PhotoLocator/AI/denoise_dncnn.py @@ -0,0 +1,41 @@ +import argparse +from pathlib import Path + +import cv2 +import torch + +from train_dncnn import DnCNN, checkpoint_model_state, denoise_image + + +def main(): + parser = argparse.ArgumentParser(description="Denoise an image using a trained PyTorch DnCNN checkpoint.") + parser.add_argument("--model", type=str, required=True, help="Path to the trained PyTorch checkpoint (.pt)") + parser.add_argument("--input", type=str, required=True, help="Path to the image to denoise") + parser.add_argument("--output", type=str, required=True, help="Path for the denoised image") + parser.add_argument("--channels", type=int, default=3, choices=[1, 3], help="Model input channels") + parser.add_argument("--tile-size", type=int, default=512, help="Maximum tile width and height") + parser.add_argument("--overlap", type=int, default=64, help="Overlap between adjacent inference tiles") + parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Inference device, e.g. cpu or cuda") + args = parser.parse_args() + + image = cv2.imread(args.input, cv2.IMREAD_UNCHANGED) + if image is None: + raise FileNotFoundError(f"Unable to read input image: {args.input}") + + model = DnCNN(channels=args.channels).to(args.device) + checkpoint = torch.load(args.model, map_location=args.device) + model.load_state_dict(checkpoint_model_state(checkpoint)) + print(f"Inference device: {args.device}") + result = denoise_image(model, image, args.channels, args.device, args.tile_size, args.overlap) + if result.shape[2] == 1: + result = result[:, :, 0] + + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + if not cv2.imwrite(str(output_path), result): + raise OSError(f"Unable to save denoised image: {output_path}") + print(f"Saved denoised image to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/PhotoLocator/AI/export_dncnn_onnx.py b/PhotoLocator/AI/export_dncnn_onnx.py new file mode 100644 index 0000000..9ffb8a6 --- /dev/null +++ b/PhotoLocator/AI/export_dncnn_onnx.py @@ -0,0 +1,57 @@ +import argparse +import os + +import onnx +import torch + +from train_dncnn import DnCNN, checkpoint_model_state + + +def main(): + parser = argparse.ArgumentParser(description="Export a trained DnCNN denoiser to a single-file embedded ONNX model.") + parser.add_argument("--model", type=str, default="dncnn_model.pt", help="Path to the saved PyTorch weights (.pt)") + parser.add_argument("--output", type=str, default="dncnn_model.onnx", help="Path for the ONNX model") + parser.add_argument("--channels", type=int, default=3, choices=[1, 3], help="Input image channels") + parser.add_argument("--height", type=int, default=256, help="Example input height") + parser.add_argument("--width", type=int, default=256, help="Example input width") + parser.add_argument("--opset", type=int, default=17, help="ONNX opset version") + args = parser.parse_args() + + model = DnCNN(channels=args.channels) + checkpoint = torch.load(args.model, map_location="cpu") + model.load_state_dict(checkpoint_model_state(checkpoint)) + model.eval() + + dummy_input = torch.randn(1, args.channels, args.height, args.width) + temp_path = args.output + ".tmp" + torch.onnx.export( + model, + (dummy_input,), + temp_path, + export_params=True, + opset_version=args.opset, + do_constant_folding=True, + input_names=["input"], + output_names=["output"], + dynamic_axes={ + "input": {0: "batch_size", 2: "height", 3: "width"}, + "output": {0: "batch_size", 2: "height", 3: "width"}, + }, + ) + + model_onnx = onnx.load(temp_path) + onnx.save_model( + model_onnx, + args.output, + save_as_external_data=False, + all_tensors_to_one_file=True, + ) + + os.remove(temp_path) + os.remove(temp_path + ".data") + + print(f"Exported ONNX model to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/PhotoLocator/AI/train_dncnn.py b/PhotoLocator/AI/train_dncnn.py new file mode 100644 index 0000000..579638d --- /dev/null +++ b/PhotoLocator/AI/train_dncnn.py @@ -0,0 +1,319 @@ +import argparse +import math +from pathlib import Path + +import cv2 +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, Dataset + + +class PairedImageDataset(Dataset): + def __init__(self, image_dir: str, clean_postfix: str = "_clean", patch_size: int = 64, channels: int = 3, patches_per_image: int = 1, use_augment: bool = True): + self.image_dir = Path(image_dir) + self.clean_postfix = clean_postfix + self.patch_size = patch_size + self.channels = channels + self.patches_per_image = max(1, patches_per_image) + self.use_augment = use_augment + + image_files = sorted(self.image_dir.glob("*")) + noisy_map = {} + clean_map = {} + + for path in image_files: + if not path.is_file(): + continue + if path.name.endswith(f"{self.clean_postfix}{path.suffix}"): + base_name = path.stem.removesuffix(self.clean_postfix) + clean_map[base_name] = path + else: + noisy_map[path.stem] = path + + unmatched_noisy_names = sorted(set(noisy_map) - set(clean_map)) + for name in unmatched_noisy_names: + print(f"No matching clean image for noisy file: {noisy_map[name].name}") + + common_names = sorted(set(noisy_map).intersection(clean_map)) + if not common_names: + raise FileNotFoundError( + f"No matching noisy/clean image pairs found in {image_dir} using clean postfix '{self.clean_postfix}'." + ) + + self.samples = [(noisy_map[name], clean_map[name]) for name in common_names] + + def __len__(self): + return len(self.samples) * self.patches_per_image + + def _load_image(self, path: Path): + #image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED) + #if image is None: + encoded = path.read_bytes() + image = cv2.imdecode(np.frombuffer(encoded, dtype=np.uint8), cv2.IMREAD_UNCHANGED) + + if image.ndim == 2: + image = image[:, :, None] + + if image.shape[2] == 1 and self.channels == 3: + image = np.repeat(image, 3, axis=2) + elif image.shape[2] == 3 and self.channels == 1: + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)[:, :, None] + + if image.shape[2] != self.channels: + raise ValueError(f"Unexpected channel count {image.shape[2]} for channel setting {self.channels}: {path}") + + image = image.astype(np.float32) / _image_max_value(image) + return image + + def __getitem__(self, index): + base_index = index // self.patches_per_image + noisy_path, clean_path = self.samples[base_index] + noisy = self._load_image(noisy_path) + clean = self._load_image(clean_path) + + if noisy.shape[:2] != clean.shape[:2]: + raise ValueError(f"Noisy and clean dimensions do not match for {noisy_path.name}: {noisy.shape[:2]} vs {clean.shape[:2]}") + + h, w = noisy.shape[:2] + if h < self.patch_size or w < self.patch_size: + raise ValueError(f"Image too small for patch size {self.patch_size}: {noisy_path}") + + if self.use_augment: + top = np.random.randint(0, h - self.patch_size + 1) + left = np.random.randint(0, w - self.patch_size + 1) + noisy = noisy[top:top + self.patch_size, left:left + self.patch_size] + clean = clean[top:top + self.patch_size, left:left + self.patch_size] + + if np.random.rand() > 0.5: + noisy = noisy[:, ::-1, :].copy() + clean = clean[:, ::-1, :].copy() + + if np.random.rand() > 0.5: + noisy = noisy[::-1, :, :].copy() + clean = clean[::-1, :, :].copy() + else: + top = (index % self.patches_per_image) * 0 + left = (index % self.patches_per_image) * 0 + noisy = noisy[0:self.patch_size, 0:self.patch_size] + clean = clean[0:self.patch_size, 0:self.patch_size] + + noisy = noisy.transpose(2, 0, 1).copy() + clean = clean.transpose(2, 0, 1).copy() + + return torch.from_numpy(noisy), torch.from_numpy(clean) + + +class DnCNN(nn.Module): + def __init__(self, channels: int = 3, num_layers: int = 17, features: int = 64): + super().__init__() + self.channels = channels + self.num_layers = num_layers + + layers = [nn.Conv2d(channels, features, kernel_size=3, stride=1, padding=1, bias=False)] + layers.append(nn.ReLU(inplace=True)) + + for _ in range(num_layers - 2): + layers.append(nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=False)) + layers.append(nn.BatchNorm2d(features)) + layers.append(nn.ReLU(inplace=True)) + + layers.append(nn.Conv2d(features, channels, kernel_size=3, stride=1, padding=1, bias=False)) + self.model = nn.Sequential(*layers) + + def forward(self, x): + residual = self.model(x) + return x - residual + + +def checkpoint_model_state(checkpoint): + if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: + return checkpoint["model_state_dict"] + return checkpoint + + +def _image_max_value(image: np.ndarray): + if image.dtype not in (np.uint8, np.uint16): + raise ValueError(f"Unsupported image dtype {image.dtype}; expected 8-bit or 16-bit unsigned pixels.") + return np.iinfo(image.dtype).max + + +def _tile_starts(length: int, tile_size: int, overlap: int): + if length <= tile_size: + return [0] + step = tile_size - overlap + starts = list(range(0, length - tile_size + 1, step)) + if starts[-1] != length - tile_size: + starts.append(length - tile_size) + return starts + + +def denoise_image(model, image: np.ndarray, channels: int, device: str, tile_size: int = 512, overlap: int = 64): + if tile_size <= 0 or overlap < 0 or overlap >= tile_size: + raise ValueError("Tile size must be positive and overlap must be between 0 and tile size.") + edge_margin = min(model.num_layers, tile_size // 2) + if overlap < 2 * edge_margin and image.ndim >= 2 and max(image.shape[:2]) > tile_size: + raise ValueError(f"Tile overlap must be at least {2 * edge_margin} pixels for this model.") + + if image.ndim == 2: + image = image[:, :, None] + output_dtype = image.dtype + image_max_value = _image_max_value(image) + if image.shape[2] == 4: + image = image[:, :, :3] + if image.shape[2] == 1 and channels == 3: + image = np.repeat(image, 3, axis=2) + elif image.shape[2] == 3 and channels == 1: + image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)[:, :, None] + if image.shape[2] != channels: + raise ValueError(f"Expected {channels} image channels, got {image.shape[2]}.") + + image = image.astype(np.float32) / image_max_value + height, width = image.shape[:2] + output = np.zeros((height, width, channels), dtype=np.float32) + weights = np.zeros((height, width, 1), dtype=np.float32) + was_training = model.training + model.eval() + try: + with torch.inference_mode(): + for top in _tile_starts(height, tile_size, overlap): + for left in _tile_starts(width, tile_size, overlap): + bottom = min(top + tile_size, height) + right = min(left + tile_size, width) + tile = image[top:bottom, left:right].transpose(2, 0, 1).copy() + tensor = torch.from_numpy(tile).unsqueeze(0).to(device) + denoised = model(tensor).squeeze(0).cpu().numpy().transpose(1, 2, 0) + tile_height, tile_width = denoised.shape[:2] + tile_weights = np.ones((tile_height, tile_width), dtype=np.float32) + if edge_margin: + ramp = np.linspace(1.0 / (edge_margin + 1), 1.0, edge_margin, dtype=np.float32) + if top > 0: + tile_weights[:edge_margin, :] *= ramp[:, None] + if bottom < height: + tile_weights[-edge_margin:, :] *= ramp[::-1, None] + if left > 0: + tile_weights[:, :edge_margin] *= ramp[None, :] + if right < width: + tile_weights[:, -edge_margin:] *= ramp[None, ::-1] + output[top:bottom, left:right] += denoised * tile_weights[:, :, None] + weights[top:bottom, left:right] += tile_weights[:, :, None] + finally: + model.train(was_training) + + output = np.clip(output / weights, 0.0, 1.0) + return np.rint(output * image_max_value).astype(output_dtype) + + +def validate_dataset(image_dir: str, clean_postfix: str): + image_path = Path(image_dir) + if not image_path.exists() or not image_path.is_dir(): + raise FileNotFoundError(f"Training dataset directory must exist: {image_dir}") + + files = sorted(image_path.glob("*")) + if not any(path.is_file() for path in files): + raise FileNotFoundError(f"Training dataset directory is empty: {image_dir}") + + matching_files = [path for path in files if path.is_file() and path.name.endswith(f"{clean_postfix}{path.suffix}")] + if not matching_files: + raise FileNotFoundError(f"No clean images with postfix '{clean_postfix}' were found in {image_dir}") + + +def main(): + parser = argparse.ArgumentParser(description="Train a DnCNN denoiser from paired noisy/clean images stored in the same folder.") + parser.add_argument("--image-dir", type=str, required=True, help="Directory containing noisy images and matching clean images with a suffix") + parser.add_argument("--clean-postfix", type=str, default="_clean", help="File suffix used by the clean image, e.g. '_clean' for 'image_clean.png'") + parser.add_argument("--output", type=str, default="dncnn_model.pt", help="Path for the trained PyTorch model") + parser.add_argument("--validation-image", type=str, help="Optional noisy image to denoise whenever a checkpoint is saved") + parser.add_argument("--resume", type=str, help="Optional checkpoint to resume from; legacy weights-only files restart the optimizer") + parser.add_argument("--validation-output-dir", type=str, help="Directory for checkpoint validation previews; defaults to a 'validation' folder beside --output") + parser.add_argument("--checkpoint-interval", type=int, default=20, help="Save a model checkpoint every N epochs (set to 0 to disable intermediate saves)") + parser.add_argument("--channels", type=int, default=3, choices=[1, 3], help="Model input channels: 1 for grayscale, 3 for RGB") + parser.add_argument("--patch-size", type=int, default=256, help="Training patch size") + parser.add_argument("--patches-per-image", type=int, default=4, help="Number of random patches sampled from each image per epoch") + parser.add_argument("--batch-size", type=int, default=8, help="Training batch size") + parser.add_argument("--epochs", type=int, default=1000, help="Total target number of epochs, including epochs in a resumed checkpoint") + parser.add_argument("--lr", type=float, default=2e-4, help="Learning rate") + parser.add_argument("--weight-decay", type=float, default=0.01, help="AdamW weight decay") + parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to train on") + args = parser.parse_args() + + if not Path(args.image_dir).exists(): + raise FileNotFoundError(f"Training dataset directory does not exist: {args.image_dir}") + validate_dataset(args.image_dir, args.clean_postfix) + + dataset = PairedImageDataset( + args.image_dir, + clean_postfix=args.clean_postfix, + patch_size=args.patch_size, + channels=args.channels, + patches_per_image=args.patches_per_image, + ) + loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=torch.cuda.is_available()) + + print(f"Image pairs: {len(dataset.samples)}") + print(f"Training device: {args.device}") + + model = DnCNN(channels=args.channels).to(args.device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + start_epoch = 0 + if args.resume: + checkpoint = torch.load(args.resume, map_location=args.device) + model.load_state_dict(checkpoint_model_state(checkpoint)) + if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: + if "optimizer_state_dict" in checkpoint: + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + start_epoch = int(checkpoint.get("epoch", 0)) + print(f"Resumed optimizer and model at epoch {start_epoch} from {args.resume}") + else: + print(f"Loaded model weights from {args.resume}; optimizer starts fresh") + criterion = nn.MSELoss() + validation_image = None + validation_output_dir = None + if args.validation_image: + validation_image = cv2.imread(args.validation_image, cv2.IMREAD_UNCHANGED) + if validation_image is None: + raise FileNotFoundError(f"Unable to read validation image: {args.validation_image}") + validation_output_dir = Path(args.validation_output_dir) if args.validation_output_dir else Path(args.output).parent / "validation" + validation_output_dir.mkdir(parents=True, exist_ok=True) + + def save_checkpoint(checkpoint_path: Path, completed_epochs: int): + torch.save({ + "epoch": completed_epochs, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + }, checkpoint_path) + print(f"Saved checkpoint to {checkpoint_path}") + if validation_image is not None: + preview = denoise_image(model, validation_image, args.channels, args.device) + preview_path = validation_output_dir / f"{checkpoint_path.stem}_validation.png" + if not cv2.imwrite(str(preview_path), preview): + raise OSError(f"Unable to save validation preview: {preview_path}") + print(f"Saved validation preview to {preview_path}") + + for epoch in range(start_epoch, args.epochs): + model.train() + running_loss = 0.0 + for noisy_batch, clean_batch in loader: + noisy_batch = noisy_batch.to(args.device) + clean_batch = clean_batch.to(args.device) + + optimizer.zero_grad() + predictions = model(noisy_batch) + loss = criterion(predictions, clean_batch) + loss.backward() + optimizer.step() + + running_loss += loss.item() * noisy_batch.size(0) + + epoch_loss = running_loss / len(dataset) + print(f"epoch={epoch + 1}/{args.epochs} loss={epoch_loss:.6f}") + + if args.checkpoint_interval > 0 and (epoch + 1) % args.checkpoint_interval == 0: + checkpoint_path = Path(f"{Path(args.output).with_suffix('')}_epoch_{epoch + 1}.pt") + save_checkpoint(checkpoint_path, epoch + 1) + + save_checkpoint(Path(args.output), max(start_epoch, args.epochs)) + + +if __name__ == "__main__": + main() diff --git a/PhotoLocator/BitmapOperations/ColorToneAdjustOperation.cs b/PhotoLocator/BitmapOperations/ColorToneAdjustOperation.cs index 308e47f..e4df540 100644 --- a/PhotoLocator/BitmapOperations/ColorToneAdjustOperation.cs +++ b/PhotoLocator/BitmapOperations/ColorToneAdjustOperation.cs @@ -151,9 +151,9 @@ public static void ColorTransformHSI2RGB(float h, float s, float i, out float r, public static void ColorTransformRGB2HSI(float r, float g, float b, out float h, out float s, out float i) { - r = Math.Clamp(r, 0f, 1f); - g = Math.Clamp(g, 0f, 1f); - b = Math.Clamp(b, 0f, 1f); + r = RealMath.Clamp(r, 0, 1); + g = RealMath.Clamp(g, 0, 1); + b = RealMath.Clamp(b, 0, 1); i = (r + g + b) / 3f; if (i == 0) { diff --git a/PhotoLocator/BitmapOperations/DnCNNDenoiseOperation.cs b/PhotoLocator/BitmapOperations/DnCNNDenoiseOperation.cs new file mode 100644 index 0000000..6b6b8c4 --- /dev/null +++ b/PhotoLocator/BitmapOperations/DnCNNDenoiseOperation.cs @@ -0,0 +1,221 @@ +using Microsoft.ML.OnnxRuntime; +using Microsoft.ML.OnnxRuntime.Tensors; +using PhotoLocator.Helpers; +using System; +using System.IO; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; + +namespace PhotoLocator.BitmapOperations; + +public class DnCNNDenoiseOperation : OperationBase +{ + public const string DefaultModelFileName = "dncnn.onnx"; + + public int TileSize { get; set; } = 512; + + public int TileOverlap { get; set; } = 32; + + public string ModelPath { get; set; } = DefaultModelFileName; + + public bool PreferDirectML { get; set; } = true; + + public bool ExpectColorModel { get; set; } + + public CancellationToken CancellationToken { get; set; } + + public Action? ProgressCallback { get; set; } + + public static string ResolveModelPath(string modelPath) + { + if (string.IsNullOrWhiteSpace(modelPath)) + modelPath = DefaultModelFileName; + if (!Path.IsPathRooted(modelPath)) + modelPath = Path.Combine(AppContext.BaseDirectory, modelPath); + return modelPath; + } + + public static string? TryFindModelInAppDirectory() + { + foreach (var candidate in new[] + { + DefaultModelFileName, + Path.Combine("models", DefaultModelFileName), + Path.Combine("ai", DefaultModelFileName) + }) + { + var fullPath = ResolveModelPath(candidate); + if (File.Exists(fullPath)) + return fullPath; + } + return null; + } + + public override void Apply() + { + if (string.IsNullOrWhiteSpace(ModelPath)) + throw new InvalidOperationException("ModelPath must be set before applying DnCNNDenoiseOperation."); + if (SrcBitmap == null) + throw new InvalidOperationException("SrcBitmap cannot be null."); + if (DstBitmap == null) + throw new InvalidOperationException("DstBitmap cannot be null."); + + var source = SrcBitmap; + if (source.PlaneCount is not 1 and not 3) + throw new NotSupportedException("DnCNNDenoiseOperation supports grayscale (1 channel) and RGB (3 channels) images only."); + + var channels = source.PlaneCount; + var height = source.Height; + var width = source.Width; + + var modelPath = ResolveModelPath(ModelPath); + if (!File.Exists(modelPath)) + throw new FileNotFoundException($"ONNX model not found: {modelPath}"); + + if (TileSize <= 0) + TileSize = 512; + if (TileOverlap < 0) + TileOverlap = 0; + if (TileOverlap >= TileSize) + TileOverlap = Math.Max(0, TileSize / 4); + + bool processGrayscaleWithColorModel = ExpectColorModel && channels == 1; + if (processGrayscaleWithColorModel) + { + source = new FloatBitmap(SrcBitmap.Width, SrcBitmap.Height, 3); + source.ProcessElementWise(SrcBitmap, (_, s) => s); + channels = 3; + } + + using var session = CreateSession(modelPath, PreferDirectML); + + if (width <= TileSize && height <= TileSize) + { + DstBitmap.New(width, height, channels); + RunTileInference(session, source, 0, 0, width, height, channels, DstBitmap, 0, 0); + return; + } + + var accumulator = new float[height, width, channels]; + var weightMap = new float[height, width]; + var step = Math.Max(1, TileSize - TileOverlap); + + for (var top = 0; top < height; top += step) + { + var tileTop = top; + var tileBottom = Math.Min(top + TileSize, height); + var tileHeight = tileBottom - tileTop; + for (var left = 0; left < width; left += step) + { + CancellationToken.ThrowIfCancellationRequested(); + var tileLeft = left; + var tileRight = Math.Min(left + TileSize, width); + var tileWidth = tileRight - tileLeft; + + var tileDst = new FloatBitmap(tileWidth, tileHeight, channels); + RunTileInference(session, source, tileLeft, tileTop, tileWidth, tileHeight, channels, tileDst, 0, 0); + var topOverlap = tileTop > 0 ? Math.Min(TileOverlap, tileHeight) : 0; + var bottomOverlap = tileBottom < height ? Math.Min(TileOverlap, tileHeight) : 0; + var leftOverlap = tileLeft > 0 ? Math.Min(TileOverlap, tileWidth) : 0; + var rightOverlap = tileRight < width ? Math.Min(TileOverlap, tileWidth) : 0; + + Parallel.For(0, tileHeight, y => + { + var globalY = tileTop + y; + var yWeight = GetRampWeight(y, tileHeight, topOverlap, bottomOverlap); + for (var x = 0; x < tileWidth; x++) + { + var globalX = tileLeft + x; + var weight = yWeight * GetRampWeight(x, tileWidth, leftOverlap, rightOverlap); + weightMap[globalY, globalX] += weight; + for (var c = 0; c < channels; c++) + accumulator[globalY, globalX, c] += tileDst.Elements[y, x * channels + c] * weight; + } + }); + ProgressCallback?.Invoke((top + (left + tileWidth) / (double)width * tileHeight) / height); + } + } + + CancellationToken.ThrowIfCancellationRequested(); + DstBitmap.New(width, height, channels); + Parallel.For(0, height, y => + { + for (var x = 0; x < width; x++) + { + if (weightMap[y, x] == 0) + throw new InvalidOperationException($"Weight map at ({x}, {y}) is zero, which should not happen."); + for (var c = 0; c < channels; c++) + DstBitmap.Elements[y, x * channels + c] = RealMath.Clamp(accumulator[y, x, c] / weightMap[y, x], 0, 1); + } + }); + + if (processGrayscaleWithColorModel) + DstBitmap.Assign(ConvertToGrayscaleOperation.ConvertToGrayscale(DstBitmap)); + } + + private static float GetRampWeight(int position, int length, int leadingOverlap, int trailingOverlap) + { + var weight = 1f; + if (leadingOverlap > 0 && position < leadingOverlap) + weight = (position + 1) / (float)(leadingOverlap + 1); + if (trailingOverlap > 0 && position >= length - trailingOverlap) + weight = Math.Min(weight, (length - position) / (float)(trailingOverlap + 1)); + return weight; + } + + private static InferenceSession CreateSession(string modelPath, bool preferDirectML) + { + if (!preferDirectML) + return new InferenceSession(modelPath); + + try + { + using var sessionOptions = new SessionOptions(); + sessionOptions.AppendExecutionProvider_DML(0); + return new InferenceSession(modelPath, sessionOptions); + } + catch (Exception exception) when (exception is OnnxRuntimeException or DllNotFoundException) + { + return new InferenceSession(modelPath); + } + } + + private static void RunTileInference(InferenceSession session, FloatBitmap source, int srcLeft, int srcTop, int tileWidth, int tileHeight, int channels, FloatBitmap dstBitmap, int dstLeft, int dstTop) + { + var inputBuffer = new float[1 * channels * tileHeight * tileWidth]; + for (var y = 0; y < tileHeight; y++) + { + var sourceY = srcTop + y; + for (var x = 0; x < tileWidth; x++) + { + var sourceX = srcLeft + x; + for (var c = 0; c < channels; c++) + { + var index = (c * tileHeight + y) * tileWidth + x; + inputBuffer[index] = source.Elements[sourceY, sourceX * channels + c]; + } + } + } + + var inputName = session.InputMetadata.Keys.First(); + using var output = session.Run([NamedOnnxValue.CreateFromTensor(inputName, new DenseTensor(inputBuffer, [1, channels, tileHeight, tileWidth]))]); + var outputBuffer = output[0].AsTensor().ToArray(); + + for (var y = 0; y < tileHeight; y++) + { + for (var x = 0; x < tileWidth; x++) + { + for (var c = 0; c < channels; c++) + { + var index = (c * tileHeight + y) * tileWidth + x; + var outValue = RealMath.Clamp(outputBuffer[index], 0, 1); + var dstY = dstTop + y; + var dstX = dstLeft + x; + if (dstY >= 0 && dstY < dstBitmap.Height && dstX >= 0 && dstX < dstBitmap.Width) + dstBitmap.Elements[dstY, dstX * channels + c] = outValue; + } + } + } + } +} diff --git a/PhotoLocator/ImageTransformCommands.cs b/PhotoLocator/ImageTransformCommands.cs index d5b1cc9..f3a21ac 100644 --- a/PhotoLocator/ImageTransformCommands.cs +++ b/PhotoLocator/ImageTransformCommands.cs @@ -294,6 +294,79 @@ await Task.Run(() => GeneralFileFormatHandler.SaveToFile(image, targetFileName, int _resizeTargetHeight = 2160; + public ICommand DenoiseCommand => new RelayCommand(async o => + { + var allSelected = _mainViewModel.GetSelectedItems(true).ToArray(); + if (allSelected.Length == 0) + return; + + var modelPath = DnCNNDenoiseOperation.TryFindModelInAppDirectory(); + if (string.IsNullOrWhiteSpace(modelPath)) + { + var dlg = new Microsoft.Win32.OpenFileDialog + { + Title = "Select DnCNN ONNX model", + Filter = "ONNX model (*.onnx)|*.onnx", + CheckFileExists = true, + CheckPathExists = true, + InitialDirectory = Path.Combine(AppContext.BaseDirectory, "AI"), + }; + if (dlg.ShowDialog() != true) + return; + modelPath = dlg.FileName; + } + + var browser = new System.Windows.Forms.FolderBrowserDialog(); + browser.InitialDirectory = Path.GetDirectoryName(allSelected[0].FullPath)!; + browser.Description = "Select target folder for denoised images"; + browser.UseDescriptionForTitle = true; + if (browser.ShowDialog() != System.Windows.Forms.DialogResult.OK) + return; + + var targetDir = browser.SelectedPath; + var targetIsSourceDir = string.Equals(targetDir, Path.GetDirectoryName(allSelected[0].FullPath), StringComparison.OrdinalIgnoreCase); + var overwriteAll = targetIsSourceDir; + + await _mainViewModel.RunProcessWithProgressBarAsync(async (progressCallback, ct) => + { + int i = 0; + foreach (var item in allSelected) + { + var targetFileName = targetIsSourceDir ? item.GetProcessedFileName() : Path.Combine(targetDir, item.Name); + if (!overwriteAll && File.Exists(targetFileName)) + { + if (MessageBox.Show(App.Current.MainWindow, $"File {targetFileName} already exists. Overwrite all conflicting files?", + "Confirm Overwrite All", MessageBoxButton.OKCancel, MessageBoxImage.Question) != MessageBoxResult.OK) + break; + overwriteAll = true; + } + + var (image, itemMetadata) = await LoadImageWithMetadataAsync(item, ct); + await Task.Run(() => + { + var sourceFloat = new FloatBitmap(image, 1); + var denoiseOperation = new DnCNNDenoiseOperation + { + SrcBitmap = sourceFloat, + DstBitmap = new FloatBitmap(sourceFloat.Width, sourceFloat.Height, sourceFloat.PlaneCount), + ModelPath = modelPath, + ExpectColorModel = true, + CancellationToken = ct, + ProgressCallback = progress => progressCallback((double)(i + progress * 0.95) / allSelected.Length), + }; + denoiseOperation.Apply(); + + var sourceBitsPerChannel = image.Format.BitsPerPixel / int.Max(1, image.Format.Masks.Count); + var resultImage = sourceBitsPerChannel > 8 ? denoiseOperation.DstBitmap.ToBitmapSource16(image.DpiX, image.DpiY, 1) : denoiseOperation.DstBitmap.ToBitmapSource(image.DpiX, image.DpiY, 1); + GeneralFileFormatHandler.SaveToFile(resultImage, targetFileName, + ExifHandler.PrepareMetadataForProcessedImage(resultImage.Format, itemMetadata), _mainViewModel.Settings); + } , ct); + item.IsChecked = false; + progressCallback((double)(++i) / allSelected.Length); + } + }, "Denoising"); + }, HasFileSelected); + public ICommand ResizeCommand => new RelayCommand(async o => { var allSelected = _mainViewModel.GetSelectedItems(true).ToArray(); diff --git a/PhotoLocator/LocalContrastView.xaml b/PhotoLocator/LocalContrastView.xaml index 917597f..27379d8 100644 --- a/PhotoLocator/LocalContrastView.xaml +++ b/PhotoLocator/LocalContrastView.xaml @@ -168,6 +168,18 @@ + + + + + + + + + + + + diff --git a/PhotoLocator/LocalContrastViewModel.cs b/PhotoLocator/LocalContrastViewModel.cs index 872b7fb..02223f5 100644 --- a/PhotoLocator/LocalContrastViewModel.cs +++ b/PhotoLocator/LocalContrastViewModel.cs @@ -22,13 +22,14 @@ namespace PhotoLocator { class LocalContrastViewModel : INotifyPropertyChanged, IImageZoomPreviewViewModel { - enum FirstParamChanged { Astro, LaplacianPyramid, LocalContrast, ColorTone, None } + enum FirstParamChanged { Astro, LaplacianPyramid, LocalContrast, Denoising, ColorTone, None } static AdjustmentValues? _adjustmentClipboard; static AdjustmentValues? _lastUsedValues; readonly DispatcherTimer _updateTimer; readonly LaplacianFilterOperation _laplacianFilterOperation = new(); readonly IncreaseLocalContrastOperation _localContrastOperation = new() { DstBitmap = new() }; + readonly DnCNNDenoiseOperation _denoiseOperation = new() { ModelPath = System.IO.Path.Combine(AppContext.BaseDirectory, "AI", "dncnn-17l-64f-2000e.onnx") }; readonly ColorToneAdjustOperation _colorToneOperation = new(); Task _previewTask = Task.CompletedTask; FirstParamChanged _firstParamChanged = FirstParamChanged.None; @@ -247,6 +248,19 @@ public double DetailHandling } = DefaultDetailHandling; public ICommand ResetDetailHandlingCommand => new RelayCommand(o => DetailHandling = DefaultDetailHandling); + public bool IsDenoisingEnabled => File.Exists(_denoiseOperation.ModelPath); + + public double DenoisingStrength + { + get; + set + { + if (SetProperty(ref field, RealMath.Clamp(value, 0, 2))) + StartUpdateTimer(FirstParamChanged.Denoising); + } + } + public ICommand ResetDenoisingCommand => new RelayCommand(o => DenoisingStrength = 0); + public ColorToneAdjustOperation.ToneAdjustment[] ToneAdjustments => _colorToneOperation.ToneAdjustments; public IEnumerable ColorTones @@ -504,6 +518,7 @@ private AdjustmentValues StoreAdjustmentValues() av.ToneMapping = ToneMapping; av.DetailHandling = DetailHandling; av.MaxStretch = MaxStretch; + av.DenoisingStrength = DenoisingStrength; av.ToneRotation = ToneRotation; if (_colorToneOperation.AreToneAdjustmentsChanged) { @@ -533,6 +548,7 @@ private void RestoreAdjustmentValues(AdjustmentValues av) ToneMapping = av.ToneMapping; DetailHandling = av.DetailHandling; MaxStretch = av.MaxStretch; + DenoisingStrength = av.DenoisingStrength; ToneRotation = av.ToneRotation; if (av.ToneAdjustments is null || av.ToneAdjustments.Length != ColorToneAdjustOperation.NumberOfHues && av.ToneAdjustments.Length != ColorToneAdjustOperation.NumberOfHues * 2) @@ -608,6 +624,29 @@ private void ApplyLocalContrastOperation() _localContrastOperation.Apply(); } + private void ApplyDenoisingOperation() + { + if (DenoisingStrength <= 0) + return; + var denoiseWeight = DenoisingStrength; + if (_denoiseOperation.DstBitmap is null) + { + _denoiseOperation.SrcBitmap ??= new(); + _denoiseOperation.SrcBitmap.New(_localContrastOperation.DstBitmap.Width, _localContrastOperation.DstBitmap.Height, _localContrastOperation.DstBitmap.PlaneCount); + _denoiseOperation.SrcBitmap.ProcessElementWise(_localContrastOperation.DstBitmap, (_, s) => (float)Math.Pow(float.Clamp(s, 0, 1), 1 / FloatBitmap.DefaultMonitorGamma)); + _denoiseOperation.DstBitmap = _denoiseOperation.SrcBitmap; + _denoiseOperation.Apply(); + if (DenoisingStrength > 1) + { + _denoiseOperation.Apply(); + denoiseWeight = 0.5 + (denoiseWeight - 1) / 2; + } + } + _localContrastOperation.DstBitmap.ProcessElementWise(_denoiseOperation.DstBitmap, (s, d) => (float)(Math.Pow(d, FloatBitmap.DefaultMonitorGamma) * denoiseWeight + s * (1 - denoiseWeight))); + if (DenoisingStrength > 1) + _denoiseOperation.DstBitmap = null; + } + private void ApplyColorToneOperation() { if (!_colorToneOperation.AreToneAdjustmentsChanged) @@ -632,6 +671,8 @@ public async Task UpdatePreviewAsync() _laplacianFilterOperation.SourceChanged(); if (_firstParamChanged < FirstParamChanged.LocalContrast) _localContrastOperation.SourceChanged(); + if (_firstParamChanged < FirstParamChanged.Denoising) + _denoiseOperation.DstBitmap = null; if (_firstParamChanged < FirstParamChanged.ColorTone) _colorToneOperation.SourceChanged(); _firstParamChanged = FirstParamChanged.None; @@ -642,6 +683,9 @@ public async Task UpdatePreviewAsync() if (SourceBitmap is null || _updateTimer.IsEnabled) return; ApplyLocalContrastOperation(); + if (SourceBitmap is null || _updateTimer.IsEnabled) + return; + ApplyDenoisingOperation(); if (SourceBitmap is null || _updateTimer.IsEnabled) return; BrightnessContrastOperation.ApplyContrast(_localContrastOperation.DstBitmap, (float)Contrast); @@ -678,7 +722,7 @@ private void SetHistogram(int[] histogram) } public bool IsNoOperation => (!IsAstroModeEnabled || AstroStretch == 0 && BackgroundRemovalSmooth == 0 && BlackPoint == 0) - && DetailHandling == 1 && ToneMapping == 1 && HighlightStrength == 0 && ShadowStrength == 0 && Contrast == DefaultContrast + && DetailHandling == 1 && ToneMapping == 1 && HighlightStrength == 0 && ShadowStrength == 0 && DenoisingStrength == 0 && Contrast == DefaultContrast && !_colorToneOperation.AreToneAdjustmentsChanged; public BitmapSource GetResultImage(bool produce16bitOutput) @@ -701,6 +745,7 @@ public BitmapSource ApplyOperations(BitmapSource source) ApplyLaplacianFilterOperation(); _localContrastOperation.SourceChanged(); ApplyLocalContrastOperation(); + ApplyDenoisingOperation(); BrightnessContrastOperation.ApplyContrast(_localContrastOperation.DstBitmap, (float)Contrast); _colorToneOperation.SourceChanged(); ApplyColorToneOperation(); diff --git a/PhotoLocator/MainWindow.xaml b/PhotoLocator/MainWindow.xaml index 0864c92..68d2e5c 100644 --- a/PhotoLocator/MainWindow.xaml +++ b/PhotoLocator/MainWindow.xaml @@ -296,6 +296,7 @@ + @@ -426,6 +427,7 @@ + diff --git a/PhotoLocator/PhotoLocator.csproj b/PhotoLocator/PhotoLocator.csproj index 921f54f..14c98c3 100644 --- a/PhotoLocator/PhotoLocator.csproj +++ b/PhotoLocator/PhotoLocator.csproj @@ -33,6 +33,7 @@ + diff --git a/PhotoLocator/PictureFileFormats/AdjustmentValues.cs b/PhotoLocator/PictureFileFormats/AdjustmentValues.cs index 2621c86..df9aae6 100644 --- a/PhotoLocator/PictureFileFormats/AdjustmentValues.cs +++ b/PhotoLocator/PictureFileFormats/AdjustmentValues.cs @@ -19,6 +19,7 @@ internal class AdjustmentValues public double HighlightStrength { get; set; } public double ShadowStrength { get; set; } public double OutlierReductionStrength { get; set; } + public double DenoisingStrength { get; set; } public double Contrast { get; set; } public double ToneMapping { get; set; } public double DetailHandling { get; set; } diff --git a/PhotoLocatorTest/BitmapOperations/DnCNNDenoiseOperationTest.cs b/PhotoLocatorTest/BitmapOperations/DnCNNDenoiseOperationTest.cs new file mode 100644 index 0000000..206db7f --- /dev/null +++ b/PhotoLocatorTest/BitmapOperations/DnCNNDenoiseOperationTest.cs @@ -0,0 +1,42 @@ +namespace PhotoLocator.BitmapOperations; + +[TestClass] +public class DnCNNDenoiseOperationTest +{ + [TestMethod] + public void Apply_ShouldRunModel_WhenPretrainedOnnxExists() + { + var modelPath = Path.Combine(AppContext.BaseDirectory, "TestData", "dncnn_test.onnx"); + if (!File.Exists(modelPath)) + Assert.Inconclusive($"Pretrained ONNX file not found: {modelPath}. Add a trained dncnn model to PhotoLocatorTest/TestData/dncnn_test.onnx to run this test."); + + var bitmap = new FloatBitmap(8, 8, 3); + for (var y = 0; y < bitmap.Height; y++) + { + for (var x = 0; x < bitmap.Width; x++) + { + bitmap.Elements[y, x * 3 + 0] = ((x * 13 + y * 7) % 17) / 16f; + bitmap.Elements[y, x * 3 + 1] = (float)(0.5 + 0.2 * Math.Sin(x * 0.7 + y * 0.9)); + bitmap.Elements[y, x * 3 + 2] = (float)(1.0 - 0.3 * Math.Cos(x * 0.5 + y * 0.4)); + } + } + + var op = new DnCNNDenoiseOperation + { + SrcBitmap = bitmap, + DstBitmap = new FloatBitmap(), + ModelPath = modelPath, + TileSize = 6, + TileOverlap = 2, + }; + + op.Apply(); + + Assert.IsNotNull(op.DstBitmap); + Assert.AreEqual(bitmap.Width, op.DstBitmap.Width); + Assert.AreEqual(bitmap.Height, op.DstBitmap.Height); + Assert.AreEqual(bitmap.PlaneCount, op.DstBitmap.PlaneCount); + Assert.IsTrue(op.DstBitmap.Elements.Cast().All(v => !float.IsNaN(v)), "Model output contains NaN values."); + Assert.IsTrue(op.DstBitmap.Elements.Cast().All(v => v is >= 0f and <= 1f), "Model output is outside the normalized [0,1] range."); + } +} diff --git a/PhotoLocatorTest/PhotoLocatorTest.csproj b/PhotoLocatorTest/PhotoLocatorTest.csproj index f228ebc..1664fe1 100644 --- a/PhotoLocatorTest/PhotoLocatorTest.csproj +++ b/PhotoLocatorTest/PhotoLocatorTest.csproj @@ -93,6 +93,9 @@ PreserveNewest + + PreserveNewest + diff --git a/PhotoLocatorTest/TestData/dncnn_test.onnx b/PhotoLocatorTest/TestData/dncnn_test.onnx new file mode 100644 index 0000000..347f545 --- /dev/null +++ b/PhotoLocatorTest/TestData/dncnn_test.onnx @@ -0,0 +1,17 @@ + PhotoLocatorTest:¹ +( +inputoutput identity_test"Identityminimal_denoise_test_graphZ7 +input. +,( +batch + +channels +height +widthb8 +output. +,( +batch + +channels +height +widthB \ No newline at end of file