Implemented ncpu arg for denoiser.

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,31 +60,32 @@ def main():
     pass
 
 @main.command()
-@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")
+@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", 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 STAR file")
-@click.option("--conf", default=0.0, help="Minimum confidence threshold from predictions")
+@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", 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.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')
+def denoise(source, project, num_workers):
     click.echo("Denoising micrographs...")
     import partinet.process_utils.pooled_denoise_proc
-    partinet.process_utils.pooled_denoise_proc.main(source,project)
+    partinet.process_utils.pooled_denoise_proc.main(source,project,num_workers)
 
 
 @main.group()

partinet/process_utils/pooled_denoise_proc.py

@@ -8,7 +8,7 @@
 import gc
 from concurrent.futures import ProcessPoolExecutor
 
-MAX_WORKERS = max(1, multiprocessing.cpu_count() // 2)
+# MAX_WORKERS = max(1, multiprocessing.cpu_count() // 2)
 
 def clahe_denoise(args):
     try:
@@ -21,7 +21,7 @@ def clahe_denoise(args):
     except Exception as e:
         logging.error(f"Failed to process {src_path}: {str(e)}")
 
-def process_directory(micrographs_dir, clahe_denoised_dir):
+def process_directory(micrographs_dir, clahe_denoised_dir,MAX_WORKERS):
     os.makedirs(clahe_denoised_dir, exist_ok=True)
     logging.info(f"Directory ready: {clahe_denoised_dir}")
 
@@ -41,25 +41,41 @@ def process_directory(micrographs_dir, clahe_denoised_dir):
             future.result()
     gc.collect()
 
-def main(source_dir, project_dir):
+def main(source_dir, project_dir, ncpu):
     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")
     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}")
+
+    max_available_cpus = multiprocessing.cpu_count()
+    max_workers = max(1, max_available_cpus // 2)
+
+    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)
+    # process_directory(source_dir, denoise_dir,ncpu)
 
 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")
+    parser.add_argument('--ncpu', type=int, default=None, help='Number of CPUs to use')
+
 
     return parser.parse_args()
 
 if __name__ == "__main__":
     args = parse_args()
-    main(args.raw, args.project)
+    main(args.raw, args.project, args.ncpu)

partinet/process_utils/split_train.py

@@ -8,8 +8,6 @@
 def main(labels_path, images_path, output_dir):
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
-    os.makedirs(os.path.join(output_dir, "images"))
-    os.makedirs(os.path.join(output_dir, "labels"))
     os.makedirs(os.path.join(output_dir, "images", "train"))
     os.makedirs(os.path.join(output_dir, "images", "val"))
     os.makedirs(os.path.join(output_dir, "labels", "train"))