Added comments, doc strings and type annotations

Author: MihinP

partinet/process_utils/guided_denoiser.py

@@ -7,91 +7,154 @@
 from numpy.fft import fft2, ifft2
 from scipy.signal import gaussian
 
-def transform(image):
+
+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
+    image = ((image - i_min) / (i_max - i_min)) * 255
     return image.astype(np.uint8)
 
 
-def standard_scaler(image):
+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
-    image = transform(image).astype(np.uint8)
-    return image
+    image = (image - mu) / sigma
+    return transform(image).astype(np.uint8)
 
-def contrast_enhancement(image):
-    enhanced_image = cv2.fastNlMeansDenoising(image, None, h=10, templateWindowSize=7, searchWindowSize=21)
-    
-    return enhanced_image
 
+def contrast_enhancement(image: np.ndarray) -> np.ndarray:
+    """
+    Enhance the contrast of the image using Non-Local Means denoising.
 
-def gaussian_kernel(kernel_size = 3):
+    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, kernel, K):
+
+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 = np.copy(img)
-    dummy = fft2(dummy)
-    kernel = fft2(kernel, s = img.shape)
+    dummy = fft2(img)
+    kernel = fft2(kernel, s=img.shape)
     kernel = np.conj(kernel) / (np.abs(kernel) ** 2 + K)
-    dummy = dummy * kernel
-    dummy = np.abs(ifft2(dummy))
-    return dummy
+    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 clahe(image):
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(16,16))
 
-    # Apply CLAHE to the image
-    img_equalized = clahe.apply(transform(image))
-    return img_equalized
+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.
 
-def guided_filter(input_image, guidance_image, radius=20, epsilon=0.1):
-    # Convert images to float32
+    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
 
-    # Compute mean values of the guidance image and input image
     mean_guidance = cv2.boxFilter(guidance_image, -1, (radius, radius))
     mean_input = cv2.boxFilter(input_image, -1, (radius, radius))
-
-    # Compute correlation and covariance of the guidance and input images
     mean_guidance_input = cv2.boxFilter(guidance_image * input_image, -1, (radius, radius))
     covariance_guidance_input = mean_guidance_input - mean_guidance * mean_input
 
-    # Compute squared mean of the guidance image
     mean_guidance_sq = cv2.boxFilter(guidance_image * guidance_image, -1, (radius, radius))
     variance_guidance = mean_guidance_sq - mean_guidance * mean_guidance
 
-    # Compute weights and mean of the weights
     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))
 
-    # Compute the filtered image
     output_image = mean_a * guidance_image + mean_b
-
     return transform(output_image)
 
 
-def denoise(image_path):
-    kernel = gaussian_kernel(kernel_size = 9)
+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)
+    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

@@ -7,25 +7,54 @@
 import argparse
 import gc
 from concurrent.futures import ProcessPoolExecutor
+from typing import List, Tuple
 
-# MAX_WORKERS = max(1, multiprocessing.cpu_count() // 2)
+# Function to perform CLAHE-based denoising
 
-def clahe_denoise(args):
+def clahe_denoise(args: Tuple[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.
+
+    Raises:
+        Exception: Logs any exceptions that occur during processing.
+    """
     try:
         src_path, dest_path = args
+        # Perform denoising
         denoised = denoise(src_path)
+        # Save the denoised image
         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)}")
 
-def process_directory(micrographs_dir, clahe_denoised_dir,MAX_WORKERS):
+# Function to process all files in a directory
+
+def process_directory(micrographs_dir: str, clahe_denoised_dir: str, max_workers: int) -> 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.
+
+    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 = []
+    tasks: List[Tuple[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)
@@ -35,47 +64,74 @@ def process_directory(micrographs_dir, clahe_denoised_dir,MAX_WORKERS):
             else:
                 tasks.append((src_path, dest_path))
 
-    with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:
+    # 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()
+        future.result()
     gc.collect()
 
-def main(source_dir, project_dir, ncpu):
+# Main function
+
+def main(source_dir: str, project_dir: str, ncpu: int) -> 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.
+
+    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())
 
-    denoise_dir = os.path.join(project_dir,"denoised")
+    # 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.")
-    
-    # num_cpus = max(multiprocessing.cpu_count(),ncpu)
-    # logging.info(f"Number of available CPUs: {num_cpus}")
-
     logging.info(f"Processing raw micrographs in {source_dir}")
     logging.info(f"Saving denoised micrographs in {denoise_dir}")
-    
-    # process_directory(source_dir, denoise_dir,ncpu)
 
-def parse_args():
+    # Process the directory
+    process_directory(source_dir, denoise_dir, ncpu)
+
+# 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.
+    """
     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')
-
-    
     return parser.parse_args()
 
+# Entry point for the script
+
 if __name__ == "__main__":
     args = parse_args()
-    main(args.raw, args.project, args.ncpu)
+    main(args.raw, args.project, args.ncpu)