Merge pull request #27 from WEHI-ResearchComputing/prepostprocess

Custom denoiser stage

Author: MihinP

Dockerfile

@@ -8,4 +8,4 @@ RUN python -m pip install --no-cache-dir --no-cache /opt/PartiNet
 
 LABEL AUTHORS Mihin Perera, Edward Yang, Julie Iskander
 LABEL MAINTAINERS Mihin Perera, Edward Yang, Julie Iskander
-LABEL VERSION v0.1.1
+LABEL VERSION v0.2.0

Singularity

@@ -18,4 +18,4 @@ from: python:3.9.19-slim-bookworm
 %labels
   AUTHORS Mihin Perera, Edward Yang, Julie Iskander
   MAINTAINERS Mihin Perera, Edward Yang, Julie Iskander
-  VERSION v0.1.1
+  VERSION v0.2.0

partinet/__init__.py

@@ -1,7 +1,7 @@
 import click
 import sys, os
 
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 
 DYNAMICDET_AVAILABLE_MODELS = ["yolov7", "yolov7x", "yolov7-w6", "yolov7-e6", "yolov7-d6", "yolov7-e6e"]
 
@@ -60,9 +60,34 @@ def main():
     pass
 
 @main.command()
-def preprocess():
-    click.echo("This will preprocess the micrographs.")
-    raise NotImplementedError("Not implemented yet!")
+@click.option("--labels", type=str, required=True, help="Path to the labels directory")
+@click.option("--images", type=str, required=True, help="Path to the images directory")
+@click.option("--output", type=str, required=True, help="Path to the output directory")
+def split(labels, images, output):
+    click.echo("Splitting micrographs for training and validation...")
+    import partinet.process_utils.split_train
+    partinet.process_utils.split_train.main(labels, images, output)
+
+@main.command()
+@click.option("--labels", type=str, required=True, help="Path to the labels directory")
+@click.option("--images", type=str, required=True, help="Path to the images directory")
+@click.option("--output", type=str, required=True, help="Path to the output STAR file")
+@click.option("--conf", type=float, default=0.0, help="Minimum confidence threshold from predictions")
+def star(labels, images, output,conf):
+    click.echo("Generating STAR file...")
+    import partinet.process_utils.star_file
+    partinet.process_utils.star_file.main(labels,images,output,conf)
+
+@main.command()
+@click.option("--source", type=str, required=True, help="Path to Raw micrographs")
+@click.option('--project', type=str, required=True, help='save denoised micrographs to project/denoised', show_default=True)
+@click.option('--num_workers', type=int, default=None, help='Number of workers for denoising micrographs')
+@click.option('--img_format', type=click.Choice(["png", "jpg", "mrc"], case_sensitive=True), default="png",show_default=True, help='Output format of denoised micrographs')
+def denoise(source, project, num_workers,img_format):
+    click.echo("Denoising micrographs...")
+    import partinet.process_utils.pooled_denoise_proc
+    partinet.process_utils.pooled_denoise_proc.main(source,project,num_workers,img_format)
+
 
 @main.group()
 def train():

partinet/process_utils/guided_denoiser.py

@@ -0,0 +1,160 @@
+# Adapted from CryoSegNet
+# https://github.com/jianlin-cheng/CryoSegNet
+
+import numpy as np
+import mrcfile
+import cv2
+from numpy.fft import fft2, ifft2
+from scipy.signal import gaussian
+
+
+def transform(image: np.ndarray) -> np.ndarray:
+    """
+    Normalize and scale an image to 8-bit range (0-255).
+
+    Args:
+        image (np.ndarray): Input image array.
+
+    Returns:
+        np.ndarray: Normalized and scaled 8-bit image.
+    """
+    i_min = image.min()
+    i_max = image.max()
+    image = ((image - i_min) / (i_max - i_min)) * 255
+    return image.astype(np.uint8)
+
+
+def standard_scaler(image: np.ndarray) -> np.ndarray:
+    """
+    Apply Gaussian blur and standardize the image to have zero mean and unit variance.
+
+    Args:
+        image (np.ndarray): Input image array.
+
+    Returns:
+        np.ndarray: Scaled and transformed image.
+    """
+    kernel_size = 9
+    image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
+    mu = np.mean(image)
+    sigma = np.std(image)
+    image = (image - mu) / sigma
+    return transform(image).astype(np.uint8)
+
+
+def contrast_enhancement(image: np.ndarray) -> np.ndarray:
+    """
+    Enhance the contrast of the image using Non-Local Means denoising.
+
+    Args:
+        image (np.ndarray): Input image array.
+
+    Returns:
+        np.ndarray: Contrast-enhanced image.
+    """
+    return cv2.fastNlMeansDenoising(image, None, h=10, templateWindowSize=7, searchWindowSize=21)
+
+
+def gaussian_kernel(kernel_size: int = 3) -> np.ndarray:
+    """
+    Generate a 2D Gaussian kernel.
+
+    Args:
+        kernel_size (int, optional): Size of the kernel. Defaults to 3.
+
+    Returns:
+        np.ndarray: Gaussian kernel.
+    """
+    h = gaussian(kernel_size, kernel_size / 3).reshape(kernel_size, 1)
+    h = np.dot(h, h.transpose())
+    h /= np.sum(h)
+    return h
+
+
+def wiener_filter(img: np.ndarray, kernel: np.ndarray, K: float) -> np.ndarray:
+    """
+    Apply Wiener filtering to the image.
+
+    Args:
+        img (np.ndarray): Input image array.
+        kernel (np.ndarray): Convolution kernel.
+        K (float): Noise-to-signal power ratio.
+
+    Returns:
+        np.ndarray: Filtered image.
+    """
+    kernel /= np.sum(kernel)
+    dummy = fft2(img)
+    kernel = fft2(kernel, s=img.shape)
+    kernel = np.conj(kernel) / (np.abs(kernel) ** 2 + K)
+    dummy *= kernel
+    return np.abs(ifft2(dummy))
+
+
+def clahe(image: np.ndarray) -> np.ndarray:
+    """
+    Apply Contrast Limited Adaptive Histogram Equalization (CLAHE) to the image.
+
+    Args:
+        image (np.ndarray): Input image array.
+
+    Returns:
+        np.ndarray: CLAHE-processed image.
+    """
+    clahe_instance = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(16, 16))
+    return clahe_instance.apply(transform(image))
+
+
+def guided_filter(input_image: np.ndarray, guidance_image: np.ndarray, radius: int = 20, epsilon: float = 0.1) -> np.ndarray:
+    """
+    Apply guided filtering to refine the image using a guidance image.
+
+    Args:
+        input_image (np.ndarray): Input image array.
+        guidance_image (np.ndarray): Guidance image array.
+        radius (int, optional): Radius of the filter. Defaults to 20.
+        epsilon (float, optional): Regularization parameter. Defaults to 0.1.
+
+    Returns:
+        np.ndarray: Filtered image.
+    """
+    input_image = input_image.astype(np.float32) / 255.0
+    guidance_image = guidance_image.astype(np.float32) / 255.0
+
+    mean_guidance = cv2.boxFilter(guidance_image, -1, (radius, radius))
+    mean_input = cv2.boxFilter(input_image, -1, (radius, radius))
+    mean_guidance_input = cv2.boxFilter(guidance_image * input_image, -1, (radius, radius))
+    covariance_guidance_input = mean_guidance_input - mean_guidance * mean_input
+
+    mean_guidance_sq = cv2.boxFilter(guidance_image * guidance_image, -1, (radius, radius))
+    variance_guidance = mean_guidance_sq - mean_guidance * mean_guidance
+
+    a = covariance_guidance_input / (variance_guidance + epsilon)
+    b = mean_input - a * mean_guidance
+    mean_a = cv2.boxFilter(a, -1, (radius, radius))
+    mean_b = cv2.boxFilter(b, -1, (radius, radius))
+
+    output_image = mean_a * guidance_image + mean_b
+    return transform(output_image)
+
+
+def denoise(image_path: str) -> np.ndarray:
+    """
+    Apply a sequence of denoising operations to an input image.
+
+    Args:
+        image_path (str): Path to the input .mrc image file.
+
+    Returns:
+        np.ndarray: Fully denoised image.
+    """
+    kernel = gaussian_kernel(kernel_size=9)
+    image = mrcfile.read(image_path)
+    image = image.T
+    image = np.rot90(image)
+    normalized_image = standard_scaler(np.array(image))
+    contrast_enhanced_image = contrast_enhancement(normalized_image)
+    weiner_filtered_image = wiener_filter(contrast_enhanced_image, kernel, K=30)
+    clahe_image = clahe(weiner_filtered_image)
+    guided_filter_image = guided_filter(clahe_image, weiner_filtered_image)
+    return guided_filter_image

partinet/process_utils/pooled_denoise_proc.py

@@ -0,0 +1,145 @@
+import os
+import cv2
+from pathlib import Path
+from partinet.process_utils.guided_denoiser import denoise
+import logging
+import multiprocessing
+import argparse
+import gc
+from concurrent.futures import ProcessPoolExecutor
+from typing import List, Tuple
+import mrcfile
+
+# Function to perform CLAHE-based denoising
+
+def clahe_denoise(args: Tuple[str, str, str]) -> None:
+    """
+    Applies the guided denoising algorithm to an input image and saves the denoised result.
+
+    Args:
+        args (Tuple[str, str]):
+            - src_path: Path to the source micrograph file.
+            - dest_path: Path to save the denoised image.
+            - img_format: output format of denoised images
+
+    Raises:
+        Exception: Logs any exceptions that occur during processing.
+    """
+    try:
+        src_path, dest_path, img_format = args
+        # Perform denoising
+        denoised = denoise(src_path)
+        if img_format == "mrc":
+            mrcfile.write(dest_path,data=denoised)
+        else:
+            cv2.imwrite(dest_path, denoised)
+        logging.info(f"Processed image {src_path} to dest. {dest_path}")
+        del denoised
+        gc.collect()
+    except Exception as e:
+        logging.error(f"Failed to process {src_path}: {str(e)}")
+
+# Function to process all files in a directory
+
+def process_directory(micrographs_dir: str, clahe_denoised_dir: str, max_workers: int, img_format: str) -> None:
+    """
+    Processes all `.mrc` files in the given directory using parallel workers for denoising.
+
+    Args:
+        micrographs_dir (str): Path to the directory containing raw micrograph files.
+        clahe_denoised_dir (str): Path to the directory where denoised images will be saved.
+        max_workers (int): Number of worker processes to use for parallel processing.
+        img_format (str): Output format of denoised images
+
+    Notes:
+        - Only `.mrc` files are processed.
+        - Existing denoised images are skipped.
+    """
+    os.makedirs(clahe_denoised_dir, exist_ok=True)
+    logging.info(f"Directory ready: {clahe_denoised_dir}")
+
+    tasks: List[Tuple[str, str, str]] = []
+    # Iterate through files in the directory
+    for file_name in os.listdir(micrographs_dir):
+        if file_name.endswith(".mrc"):
+            src_path = os.path.join(micrographs_dir, file_name)
+            dest_path = os.path.join(clahe_denoised_dir, file_name.replace(".mrc", "."+img_format))
+            if os.path.exists(dest_path):
+                logging.info(f"{dest_path} already exists!")
+            else:
+                tasks.append((src_path, dest_path, img_format))
+
+    # Parallel processing of tasks
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(clahe_denoise, task) for task in tasks]
+    for future in futures:
+        future.result()
+    gc.collect()
+
+# Main function
+
+def main(source_dir: str, project_dir: str, ncpu: int, img_format: str) -> None:
+    """
+    Main function to set up logging, determine available CPUs, and start the denoising process.
+
+    Args:
+        source_dir (str): Path to the directory containing raw micrographs.
+        project_dir (str): Path to the project directory where denoised images will be saved.
+        ncpu (int): Number of CPUs to use for parallel processing. Defaults to half the available CPUs.
+        img_format (str): Output format of denoised images
+
+    Notes:
+        - Logging is configured to write messages to a log file and the console.
+        - Ensures at least one CPU is used for processing.
+    """
+    # Configure logging
+    logging.basicConfig(filename="partinet_denoise.log", level=logging.INFO, format='%(asctime)s - %(message)s')
+    logging.getLogger().addHandler(logging.StreamHandler())
+
+    # Prepare output directory and log file
+    denoise_dir = os.path.join(project_dir, "denoised")
+    logger_name = project_dir + "/partinet_denoise.log"
+    logging.basicConfig(filename=logger_name, level=logging.INFO, format='%(asctime)s - %(message)s')
+
+    # Determine number of available CPUs
+    max_available_cpus = multiprocessing.cpu_count()
+    max_workers = max(1, max_available_cpus // 2)
+
+    # Adjust number of CPUs to use
+    if ncpu is not None:
+        ncpu = min(ncpu, max_workers)
+    else:
+        ncpu = max_workers
+
+    logging.info(f"Using {ncpu} workers out of {max_available_cpus} available CPUs.")
+    logging.info(f"Processing raw micrographs in {source_dir}")
+    logging.info(f"Saving denoised micrographs in {denoise_dir}")
+
+    # Process the directory
+    process_directory(source_dir, denoise_dir, ncpu, img_format)
+
+# Command-line argument parsing
+
+def parse_args() -> argparse.Namespace:
+    """
+    Parses command-line arguments for the script.
+
+    Returns:
+        argparse.Namespace: Parsed arguments with the following attributes:
+            - raw: Path to raw micrographs.
+            - project: Path to the project directory.
+            - ncpu: Number of CPUs to use.
+            - img_format: output format of denoised images
+    """
+    parser = argparse.ArgumentParser(description="Denoise micrographs with guided CryoSegNet-style filter")
+    parser.add_argument("--raw", required=True, help="Path to raw micrographs")
+    parser.add_argument("--project", required=True, help="Denoised micrographs saved in project/denoised")
+    parser.add_argument('--ncpu', type=int, default=None, help='Number of CPUs to use')
+    parser.add_argument('--img_format', type=str, choices=["png","jpg","mrc"], default="png", help='Output format of denoised images')
+    return parser.parse_args()
+
+# Entry point for the script
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args.raw, args.project, args.ncpu, args.img_format)