[WIP]: Add medical image denoising demo for anisotropic diffusion

python/demo/medical_imaging/data/download_cbis_ddsm.py

@@ -0,0 +1,112 @@
+import os
+import numpy as np
+import zipfile
+from PIL import Image
+from pathlib import Path
+import argparse
+import shutil
+
+def download_dataset(output_dir="data/cbis_ddsm"):
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    os.system(f"kaggle datasets download -d awsaf49/cbis-ddsm-breast-cancer-image-dataset -p {output_dir}")
+
+    zip_file = output_path / "cbis-ddsm-breast-cancer-image-dataset.zip"
+    if zip_file.exists():
+        print("Extracting dataset...")
+        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+            zip_ref.extractall(output_path)
+        zip_file.unlink()
+        return True
+    else:
+        print("Error: Dataset download failed.")
+        return False
+
+def prepare_sample_images(dataset_dir="data/cbis_ddsm", sample_dir="data/samples", n_samples=5):
+    sample_path = Path(sample_dir)
+    sample_path.mkdir(parents=True, exist_ok=True)
+
+    dataset_path = Path(dataset_dir)
+    image_files = list(dataset_path.rglob("*.png"))[:n_samples]
+
+    if not image_files:
+        print("No images found. Please download the dataset first.")
+        return
+
+    print(f"Copying {len(image_files)} sample images...")
+
+    for i, img_file in enumerate(image_files):
+        img = Image.open(img_file).convert('L')
+        max_size = 1024
+        if max(img.size) > max_size:
+            img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+        output_file = sample_path / f"sample_{i:02d}.png"
+        img.save(output_file)
+        print(f"  Saved: {output_file}")
+
+    print(f"Sample images saved to {sample_dir}/")
+
+def create_synthetic_test_image(output_file="data/synthetic_test.png", size=(512, 512)):
+    # Create clean image with geometric shapes
+    img = np.zeros(size, dtype=np.float64)
+
+    y, x = np.ogrid[:size[0], :size[1]]
+
+    # Circle 1 (top-left) - bright white
+    cx1, cy1, r1 = size[1]//4, size[0]//4, 60
+    mask1 = (x - cx1)**2 + (y - cy1)**2 <= r1**2
+    img[mask1] = 1.0
+
+    # Circle 2 (bottom-right) - slightly dimmer
+    cx2, cy2, r2 = 3*size[1]//4, 3*size[0]//4, 80
+    mask2 = (x - cx2)**2 + (y - cy2)**2 <= r2**2
+    img[mask2] = 0.9
+
+    # Rectangle (center)
+    img[size[0]//3:2*size[0]//3, size[1]//2-40:size[1]//2+40] = 0.8
+
+    # Add light Gaussian noise
+    np.random.seed(42)  # Reproducible
+    noise = np.random.normal(0, 0.03, size)
+    noisy_img = img + noise
+
+    # Clip to [0, 1]
+    noisy_img = np.clip(noisy_img, 0, 1)
+
+    # Convert to 8-bit with high quality
+    img_uint8 = (noisy_img * 255).astype(np.uint8)
+
+    # Save with maximum quality
+    Image.fromarray(img_uint8).save(output_file, quality=100, optimize=False)
+
+    print(f"Improved synthetic test image saved to {output_file}")
+    print(f"    Image size: {size}")
+    print(f"    Normalized range: [0, 1]")
+    print(f"    Noise level: LIGHT (std=0.03)")
+    print(f"    Features: 2 circles + 1 rectangle")
+    print(f"    Quality: HIGH (no compression)")
+
+    return output_file
+
+if __name__ == "__main__":   
+    parser = argparse.ArgumentParser(description="Download and prepare CBIS-DDSM dataset")
+    parser.add_argument("--download", action="store_true", help="Download full dataset")
+    parser.add_argument("--samples", action="store_true", help="Create sample subset")
+    parser.add_argument("--synthetic", action="store_true", help="Create synthetic test image")
+    parser.add_argument("--all", action="store_true", help="Do all of the above")
+
+    args = parser.parse_args()
+
+    if args.all or args.synthetic:
+        create_synthetic_test_image()
+
+    if args.all or args.download:
+        success = download_dataset()
+        if success and (args.all or args.samples):
+            prepare_sample_images()
+    elif args.samples:
+        prepare_sample_images()
+
+    if not any([args.download, args.samples, args.synthetic, args.all]):
+        parser.print_help()

python/demo/medical_imaging/demo/demo_anisotropic_diffusion_medical.py

@@ -0,0 +1,205 @@
+import numpy as np
+from mpi4py import MPI
+from dolfinx import mesh, fem, io
+from dolfinx.fem.petsc import assemble_matrix, assemble_vector, LinearProblem
+import ufl
+from ufl import grad, div, dx, dot, sqrt, inner
+from petsc4py import PETSc
+import sys
+import os
+from basix.ufl import element
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+from utils.image_processing import (
+    load_image, save_image, create_image_mesh, 
+    image_to_function, function_to_image,
+    compute_psnr, compute_ssim, edge_preservation_index
+)
+from data.download_cbis_ddsm import create_synthetic_test_image
+
+class AnisotropicDiffusion:
+    def __init__(self, image_shape, kappa=30.0, dt=0.1, comm=MPI.COMM_WORLD):
+        self.image_shape = image_shape
+        self.kappa = kappa
+        self.dt = dt
+        self.comm = comm
+
+        self.domain = create_image_mesh(image_shape, comm)
+
+        P1 = element("Lagrange", self.domain.basix_cell(), 1)
+        self.V = fem.functionspace(self.domain, P1)
+
+        self.u = ufl.TrialFunction(self.V)
+        self.v = ufl.TestFunction(self.V)
+
+        self.u_n = fem.Function(self.V)
+
+    def diffusion_coefficient(self, grad_u):
+        # Compute gradient magnitude with small epsilon to avoid division by zero
+        grad_magnitude = sqrt(dot(grad_u, grad_u) + 1e-10)
+
+        # Perona-Malik diffusion coefficient
+        c = 1.0 / (1.0 + (grad_magnitude / self.kappa)**2)
+
+        return c
+
+    def setup_variational_problem(self):
+        c = self.diffusion_coefficient(grad(self.u_n))
+
+        F = (self.u - self.u_n) * self.v * dx + \
+            self.dt * c * dot(grad(self.u), grad(self.v)) * dx
+
+        a = ufl.lhs(F)
+        L = ufl.rhs(F)
+
+        return a, L
+
+    def solve_timestep(self, a, L):
+        A = assemble_matrix(fem.form(a))
+        A.assemble()
+        b = assemble_vector(fem.form(L))
+        b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+
+        solver = PETSc.KSP().create(self.comm)
+        solver.setOperators(A)
+        solver.setType(PETSc.KSP.Type.CG)
+
+        pc = solver.getPC()
+        pc.setType(PETSc.PC.Type.JACOBI)
+
+        u_new = fem.Function(self.V)
+        solver.solve(b, u_new.x.petsc_vec)
+        u_new.x.scatter_forward()
+
+        return u_new
+
+    def denoise(self, image, n_iterations=20, verbose=True):
+        self.u_n = image_to_function(image, self.V)
+
+        a, L = self.setup_variational_problem()
+
+        for iteration in range(n_iterations):
+            if verbose and self.comm.rank == 0:
+                print(f"Iteration {iteration + 1}/{n_iterations}")
+
+            u_new = self.solve_timestep(a, L)
+
+            self.u_n.x.array[:] = u_new.x.array[:]
+
+            a, L = self.setup_variational_problem()
+
+        denoised = function_to_image(self.u_n, self.image_shape)
+
+        return denoised
+
+def demo_synthetic_image():
+    print(f"Demo: Anisotropic Diffusion on Synthetic Image\n")
+
+    # Load or create synthetic test image
+    test_image_path = "data/synthetic_test.png"
+
+    if not os.path.exists(test_image_path):
+        print("Creating synthetic test image...")
+        os.makedirs("data", exist_ok=True)
+        create_synthetic_test_image(test_image_path)
+
+    # Load image
+    print(f"Loading image: {test_image_path}")
+    image = load_image(test_image_path, normalize=True)
+    print(f"Image shape: {image.shape}")
+
+    solver = AnisotropicDiffusion(
+        image_shape=image.shape,
+        kappa=20.0,
+        dt=0.05
+    )
+
+    print("\nApplying anisotropic diffusion...")
+    denoised = solver.denoise(image, n_iterations=40, verbose=True)
+
+    # Compute quality metrics
+    print(f"Quality Metrics:\n")
+    psnr = compute_psnr(image, denoised)
+    ssim = compute_ssim(image, denoised)
+    epi = edge_preservation_index(image, denoised)
+
+    print(f"PSNR: {psnr:.2f} dB")
+    print(f"SSIM: {ssim:.4f}")
+    print(f"Edge Preservation Index: {epi:.4f}")
+
+    # Save results
+    os.makedirs("output", exist_ok=True)
+    save_image(denoised, "output/synthetic_denoised.png")
+    print(f"\nDenoised image saved to: output/synthetic_denoised.png")
+
+    return image, denoised
+
+
+def demo_medical_image(image_path):
+    print(f"Demo: Anisotropic Diffusion on Medical Image\n")
+
+    # Load image
+    print(f"Loading image: {image_path}")
+
+    target_size = (512, 512)
+    image = load_image(image_path, normalize=True, target_size=target_size)
+    print(f"Image shape: {image.shape}")
+
+    print("\nInitializing anisotropic diffusion solver...")
+    solver = AnisotropicDiffusion(
+        image_shape=image.shape,
+        kappa=20.0,
+        dt=0.05
+    )
+
+    print("\nApplying anisotropic diffusion...")
+    denoised = solver.denoise(image, n_iterations=30, verbose=True)
+
+    print(f"Quality Metrics:\n")
+    psnr = compute_psnr(image, denoised)
+    ssim = compute_ssim(image, denoised)
+    epi = edge_preservation_index(image, denoised)
+
+    print(f"PSNR: {psnr:.2f} dB")
+    print(f"SSIM: {ssim:.4f}")
+    print(f"Edge Preservation Index: {epi:.4f}")
+
+    # Save results
+    os.makedirs("output", exist_ok=True)
+    output_path = "output/medical_denoised.png"
+    save_image(denoised, output_path)
+    print(f"\nDenoised image saved to: {output_path}")
+
+    return image, denoised
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Anisotropic diffusion for medical image denoising"
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        help="Path to input image (if not provided, uses synthetic test image)"
+    )
+    parser.add_argument(
+        "--kappa",
+        type=float,
+        default=30.0,
+        help="Gradient threshold (default: 30.0)"
+    )
+    parser.add_argument(
+        "--iterations",
+        type=int,
+        default=20,
+        help="Number of iterations (default: 20)"
+    )
+
+    args = parser.parse_args()
+
+    if args.image:
+        demo_medical_image(args.image)
+    else:
+        demo_synthetic_image()

python/demo/medical_imaging/requirements.txt

@@ -0,0 +1,24 @@
+# Core Dependencies 
+# (Based on Last Stable Release: Dolfinx v0.9.0)
+fenics-dolfinx==0.9.0
+numpy>=1.24.0
+matplotlib>=3.7.0
+scipy>=1.10.0
+
+# Image Processing
+pillow>=10.0.0
+scikit-image>=0.21.0
+
+# Visualization
+pyvista>=0.43.0
+
+# Data handling
+pandas>=2.0.0
+kaggle>=1.6.0 # Unique to the CBIS DDSM Dataset
+
+# Testing
+pytest>=7.4.0
+pytest-mpi>=0.6
+
+# For DICOM
+pydicom>=2.4.0