feat: enhance CLI with additional options for processing and tracking, refactor some part of the code in utils

Author: loic-lb

src/prismtoolbox/cli/main.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+from typing import Annotated
+
 import sys
 import typer
 
@@ -23,15 +26,24 @@
 app.add_typer(app_preprocessing, name="preprocessing", help="Preprocessing commands of the PrismToolBox.")
 
 @app.callback()
-def main(verbose: int = typer.Option(0, "--verbose", "-v", count=True, help="Increase verbosity")):
+def main(
+    ctx: typer.Context,
+    verbose: Annotated[int, typer.Option("-v", count=True, help="Increase verbosity.", show_default=False)] = 0,
+    skip_errors: Annotated[bool, typer.Option(help="Skip slides which raised an error during processing.")] = False,
+    skip_existing: Annotated[bool, typer.Option(help="Skip existing files.")] = True,
+    tracking_file: Annotated[bool, typer.Option(help="Enable tracking file for processing status.")] = False,
+):
     if verbose >= 2:
         level = logging.DEBUG
     elif verbose == 1:
         level = logging.INFO
     else:
         level = logging.WARNING
     logging.basicConfig(level=level)
-
+    ctx.ensure_object(dict)
+    ctx.obj["skip_errors"] = skip_errors
+    ctx.obj["skip_existing"] = skip_existing
+    ctx.obj["tracking_file"] = tracking_file
 
 @app.command()
 def version():

src/prismtoolbox/cli/preprocessing.py

@@ -1,12 +1,22 @@
+from typing import Annotated
+
 import os
 import typer
 import logging
+import pandas as pd
 
-from enum import Enum
-from typing import Annotated
 from pathlib import Path
 
-from .utils import load_config_file
+from .utils import (
+    Engine,
+    ContoursExtension,
+    PatchExtension,
+    PatchMode,
+    load_config_file,
+    update_tracking_df,
+    extract_contours_from_slide,
+    extract_patches_from_slide,
+)
 
 log = logging.getLogger(__name__)
 
@@ -16,26 +26,10 @@
     no_args_is_help=True,
     add_completion=True,
     rich_markup_mode="rich")
-
-class Engine(str, Enum):
-    openslide = "openslide"
-    tiffslide = "tiffslide"
-    
-class ContoursExtension(str, Enum):
-    geojson = "geojson"
-    pickle = "pickle"
-    
-class PatchExtension(str, Enum):
-    h5 = "h5"
-    geojson = "geojson"
-    
-class PatchMode(str, Enum):
-    contours = "contours"
-    roi = "roi"
-    all = "all"
-    
+        
 @app_preprocessing.command(no_args_is_help=True)
 def contour(
+    ctx: typer.Context,
     slide_directory: Annotated[str, typer.Argument(
         help="Path to the directory containing the files.")],
     results_directory: Annotated[str, typer.Argument(
@@ -46,19 +40,18 @@ def contour(
     annotations_directory: Annotated[str | None, typer.Option(
         help="Path to the directory containing the annotations."
     )] = None,
-    contours_exts: Annotated[list[ContoursExtension], typer.Option(
-        help="File extension for the contours annotations.",
+    contour_exts: Annotated[list[ContoursExtension], typer.Option(
+        help="File extensions for the contours annotations.",
         case_sensitive=False,
     )] = [ContoursExtension.pickle],
-    config_file: Annotated[str, typer.Option(
+    config_file: Annotated[str | None, typer.Option(
         help="Path to the configuration file for tissue extraction."
     )] = None,
     visualize: Annotated[bool, typer.Option(
         help="Visualize the contours extracted.",
     )] = False,
 ):
     """Extract tissue contours from the slides in a specified directory."""
-    import prismtoolbox as ptb
 
     # Set default parameters
     params_detect_tissue = {
@@ -75,7 +68,7 @@ def contour(
         "line_thickness": 50,
     }
 
-    if not os.path.exists(config_file):
+    if config_file is None or not os.path.exists(config_file):
         log.info(f"Using default parameters for tissue extraction.")
     else:
         log.info(f"Using parameters from config file: {config_file}")
@@ -95,31 +88,62 @@ def contour(
         directory_visualize = os.path.join(results_directory, f"contoured_images")
         Path(directory_visualize).mkdir(parents=True, exist_ok=True)
 
+    if ctx.obj.get("tracking_file", True):
+        # Create an empty tracking file to keep track of the processing status
+        tracking_file_path = os.path.join(results_directory, "tracking_contouring.csv")
+        if not os.path.exists(tracking_file_path):
+            # Create the tracking file if it does not exist
+            col_names = ["file_name", "nb contours"] + list(params_detect_tissue.keys()) + ["status", "error message", "timestamp"]
+            tracking_df = pd.DataFrame(columns=col_names)
+            tracking_df.to_csv(tracking_file_path, index=False)
+            
     # Iterate over the files in the directory
     for file_name in os.listdir(slide_directory):
-        # Load the image
-        WSI_object = ptb.WSI(os.path.join(slide_directory, file_name), engine=engine)
-        print(f"Processing {WSI_object.slide_name}...")
-
-        if f"{WSI_object.slide_name}.pkl" in os.listdir(directory_contours):
-            continue
-        
-        # Extract the contours from the image
-        WSI_object.detect_tissue(**params_detect_tissue)
-        # Apply pathologist annotations
-        if annotations_directory is not None:
-            WSI_object.apply_pathologist_annotations(os.path.join(annotations_directory, f"{WSI_object.slide_name}.geojson"))
-        # Save extracted contours
-        for contours_ext in contours_exts:
-            WSI_object.save_tissue_contours(directory_contours, file_format=contours_ext)
-        if visualize:
-            # Visualize the extracted contours on the tissue
-            img = WSI_object.visualize(**params_visualize_WSI)
-            img.save(os.path.join(directory_visualize, f"{WSI_object.slide_name}.jpg"))
+        slide_path = os.path.join(slide_directory, file_name)
+        try:
+            already_processed, nb_contours = extract_contours_from_slide(
+                slide_path,
+                engine,
+                directory_contours,
+                params_detect_tissue,
+                contour_exts,
+                annotations_directory,
+                visualize,
+                directory_visualize,
+                params_visualize_WSI,
+                ctx.obj["skip_existing"]
+            )
+            if ctx.obj.get("tracking_file", True):
+                # Update the tracking file with the processing status
+                update_tracking_df(
+                    tracking_file_path,
+                    file_name,
+                    ("contours", nb_contours),
+                    params_dict=params_detect_tissue,
+                    already_processed=already_processed,
+                )
+        except Exception as e:
+            if ctx.obj.get("tracking_file", True):
+                # Update the tracking file with the processing status
+                update_tracking_df(
+                    tracking_file_path,
+                    file_name,
+                    nb_objects=("contours", 0),
+                    params_dict=params_detect_tissue,
+                    error_message=str(e)
+                )
+            if ctx.obj.get("skip_errors", True):
+                log.warning(f"Skipping slide {file_name}: {e}")
+                continue
+            else:
+                log.error(f"Error processing slide {file_name}:")
+                raise
+       
     print("Contours extracted and saved successfully.")
 
 @app_preprocessing.command(no_args_is_help=True)
 def patchify(
+    ctx: typer.Context,
     slide_directory: Annotated[str, typer.Argument(
         help="Path to the directory containing the files.")],
     results_directory: Annotated[str, typer.Argument(
@@ -134,34 +158,36 @@ def patchify(
         help="Engine to use for reading the slides.",
         case_sensitive=False)] = Engine.openslide,
     patch_exts: Annotated[list[PatchExtension], typer.Option(
-        help="File extension for the patches.",
+        help="File extensions for the patches.",
         case_sensitive=False,
     )] = [PatchExtension.h5],
     mode: Annotated[PatchMode, typer.Option(
         help="The mode to use for patch extraction. Possible values are 'contours', 'roi', and 'all'.",
         case_sensitive=False,
     )] = PatchMode.all,
-    config_file: Annotated[str, typer.Option(
+    config_file: Annotated[str | None, typer.Option(
         help="Path to the configuration file for patch extraction."
     )] = None,
     stitch: Annotated[bool, typer.Option(
         help="Whether to stitch the extracted patches into a single image for visualization.",
     )] = True,
+    force_patch_extraction: Annotated[bool, typer.Option(
+        help="Force patch extraction using mode 'all', even if no contours were detected when"
+        "using 'contours' mode (useful if problem to extract contours on some slides).",
+    )] = False,
+    num_workers: Annotated[int, typer.Option(
+        help="Number of workers to use for parallel processing.",
+        min=1,
+    )] = 10,
 ):
     """Extract patches from the slides in a specified directory."""
-    import prismtoolbox as ptb
-    assert mode == "contours" and contours_directory is not None, \
-        "If the mode is 'contours', you must provide a directory with contours annotations. " \
-        "Please use the `contour` command to extract contours first."
-    assert mode == "roi" and roi_csv is not None, \
-        "If the mode is 'roi', you must provide a file with the ROI coordinates."
-        
     # Set default parameters
     params_patches = {"patch_size": 256, "patch_level": 0, "overlap": 0,
-                      "units": ["px", "px"], "contours_mode": "four_pt", "rgb_threshs": [2, 240], "percentages": [0.6, 0.9]}
+                      "units": ["px", "px"], "contours_mode": "four_pt", "rgb_threshs": [2, 240], 
+                      "percentages": [0.6, 0.9]}
     params_stitch_WSI = {"vis_level": 4, "draw_grid": False}
 
-    if not os.path.exists(config_file):
+    if config_file is None or not os.path.exists(config_file):
         log.info(f"Using default parameters for tissue extraction.")
     else:
         log.info(f"Using parameters from config file: {config_file}")
@@ -173,7 +199,7 @@ def patchify(
             params_stitch_WSI = load_config_file(config_file,
                                                  dict_to_update=params_stitch_WSI,
                                                  key_to_check='stitching_settings')
-                        
+    params_patches = {**{"mode": mode}, **params_patches}
     # Path to the directory where the patches will be saved
     directory_patches = os.path.join(results_directory,
                                      f"patches_{params_patches['patch_size']}_overlap"
@@ -186,34 +212,62 @@ def patchify(
                                         f"stitched_images_{params_patches['patch_size']}_overlap"
                                         f"_{params_patches['overlap']}")
         Path(directory_stitch).mkdir(parents=True, exist_ok=True)
-        
+    
+    if ctx.obj.get("tracking_file", True):
+        # Create an empty tracking file to keep track of the processing status
+        tracking_file_path = os.path.join(results_directory, "tracking_patching.csv")
+        if not os.path.exists(tracking_file_path):
+            # Create the tracking file if it does not exist
+            col_names = ["file_name", "nb patches"] + list(params_patches.keys()) + ["status", "error message", "timestamp"]
+            tracking_df = pd.DataFrame(columns=col_names)
+            tracking_df.to_csv(tracking_file_path, index=False)
+            
     # Iterate over the files in the directory
     for file_name in os.listdir(slide_directory):
-        # Load the image
-        WSI_object = ptb.WSI(os.path.join(slide_directory, file_name), engine=engine)
-        print(f"Processing {WSI_object.slide_name}...")
-
-        if f"{WSI_object.slide_name}.h5" in os.listdir(directory_patches):
-            continue
-
-        if mode == "roi":
-            # Set the region of interest for the image
-            WSI_object.set_roi(rois_df_path=roi_csv)
-        
-        elif mode == "contours":
-            # Load the contours for the image
-            WSI_object.load_tissue_contours(os.path.join(contours_directory, f"{WSI_object.slide_name}.pkl"))
-
-        # Extract patches from the contours
-        WSI_object.extract_patches(mode=mode, **params_patches)
-        # Save the extracted patches
-        for patch_ext in patch_exts:
-            WSI_object.save_patches(directory_patches, file_format=patch_ext)
-        if stitch:
-            # Stitch the extracted patches
-            if params_stitch_WSI["vis_level"] >= len(WSI_object.level_dimensions):
+        slide_path = os.path.join(slide_directory, file_name)
+        try:
+            already_processed, nb_patches = extract_patches_from_slide(
+                slide_path,
+                engine,
+                directory_patches,
+                params_patches,
+                patch_exts,
+                roi_csv,
+                contours_directory,
+                stitch,
+                directory_stitch,
+                params_stitch_WSI,
+                force_patch_extraction,
+                ctx.obj["skip_existing"],
+                num_workers=num_workers
+            )
+            if ctx.obj.get("tracking_file", True):
+                # Update the tracking file with the processing status
+                update_tracking_df(
+                    tracking_file_path,
+                    file_name,
+                    ("patches", nb_patches),
+                    params_dict=params_patches,
+                    already_processed=already_processed
+                )
+        except Exception as e:
+            if ctx.obj.get("tracking_file", True):
+                # Update the tracking file with the processing status
+                update_tracking_df(
+                    tracking_file_path,
+                    file_name,
+                    ("patches", 0),
+                    params_dict=params_patches,
+                    error_message=str(e)
+                )
+            if ctx.obj.get("skip_errors", True):
+                log.warning(f"Skipping slide {file_name}: {e}")
                 continue
-            img = WSI_object.stitch(**params_stitch_WSI)
-
-            img.save(os.path.join(directory_stitch, f"{WSI_object.slide_name}.jpg"))
-    print("Done !")
+            else:
+                log.error(f"Error processing slide {file_name}:")
+                raise
+        if params_patches["mode"] != mode:
+            log.warning(f"Mode {params_patches['mode']} was used for patch extraction instead of {mode}." \
+                "Reverting back to mode {mode} for the next slide.")
+            params_patches["mode"] = mode
+    print("Patches extracted and saved successfully.")