Added cryosegnet-style guided denoise filter command with multiprocessing layer.

Author: MihinP

partinet/__init__.py

@@ -61,8 +61,8 @@ def main():
 @click.option("--labels", required=True, help="Path to the labels directory")
 @click.option("--images", required=True, help="Path to the images directory")
 @click.option("--output", required=True, help="Path to the output directory")
-def preprocess(labels, images, output):
-    click.echo("Preprocessing the micrographs...")
+def train_split(labels, images, output):
+    click.echo("Splitting micrographs for training and validation...")
     import partinet.split_train
     partinet.split_train.main(labels, images, output)
 
@@ -76,6 +76,14 @@ def star(labels, images, output,conf):
     import partinet.star_file
     partinet.star_file.main(labels,images,output,conf)
 
+@main.command()
+@click.option("--source", required=True, help="Path to Raw micrographs")
+@click.option('--project', required=True, help='save denoised micrographs to project/denoised', show_default=True))
+def denoise(source, project):
+    click.echo("Denoising micrographs...")
+    import partinet.pooled_denoise_proc
+    partinet.pooled_denoise_proc(source,project)
+
 
 @main.group()
 def train():

partinet/guided_denoiser.py

@@ -0,0 +1,97 @@
+# 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):
+    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):
+    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
+
+def contrast_enhancement(image):
+    enhanced_image = cv2.fastNlMeansDenoising(image, None, h=10, templateWindowSize=7, searchWindowSize=21)
+    
+    return enhanced_image
+
+
+def gaussian_kernel(kernel_size = 3):
+    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):
+    kernel /= np.sum(kernel)
+    dummy = np.copy(img)
+    dummy = fft2(dummy)
+    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
+
+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, guidance_image, radius=20, epsilon=0.1):
+    # Convert images to float32
+    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)
+    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/pooled_denoise_proc.py

@@ -0,0 +1,65 @@
+import os
+import cv2
+from pathlib import Path
+from guided_denoiser import denoise
+import logging
+import multiprocessing
+import argparse
+import gc
+from concurrent.futures import ProcessPoolExecutor
+
+MAX_WORKERS = max(1, multiprocessing.cpu_count() // 2)
+
+def clahe_denoise(args):
+    try:
+        src_path, dest_path = args
+        denoised = denoise(src_path)
+        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):
+    os.makedirs(clahe_denoised_dir, exist_ok=True)
+    logging.info(f"Directory ready: {clahe_denoised_dir}")
+
+    tasks = []
+    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", ".png"))
+            if os.path.exists(dest_path):
+                logging.info(f"{dest_path} already exists!")
+            else:
+                tasks.append((src_path, dest_path))
+
+    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()
+
+def main(source_dir, project_dir):
+    logging.basicConfig(filename="partinet_denoise.log", level=logging.INFO, format='%(asctime)s - %(message)s')
+    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')
+    num_cpus = multiprocessing.cpu_count()
+    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)
+
+def parse_args():
+    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")
+    
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args.raw, args.project)