napari-tmidas 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

napari_tmidas/processing_functions/careamics_denoising.py

@@ -0,0 +1,324 @@
+# processing_functions/careamics_denoising.py
+"""
+Processing functions for denoising images using CAREamics.
+
+This module provides functionality to denoise images using various models from CAREamics,
+including Noise2Void (N2V) and CARE models. The functions support both 2D and 3D data.
+
+The functions will automatically create and manage a dedicated environment for CAREamics
+if it's not already installed in the main environment.
+"""
+import os
+
+import numpy as np
+
+from napari_tmidas._registry import BatchProcessingRegistry
+
+# Import the environment manager for CAREamics
+from napari_tmidas.processing_functions.careamics_env_manager import (
+    create_careamics_env,
+    is_env_created,
+    run_careamics_in_env,
+)
+
+# Check if CAREamics is directly available in current environment
+try:
+    from careamics import CAREamist
+
+    CAREAMICS_AVAILABLE = True
+    USE_DEDICATED_ENV = False
+    print("CAREamics found in current environment, using direct import")
+except ImportError:
+    CAREAMICS_AVAILABLE = False
+    USE_DEDICATED_ENV = True
+    print(
+        "CAREamics not found in current environment, will use dedicated environment"
+    )
+
+
+@BatchProcessingRegistry.register(
+    name="CAREamics Denoise (N2V/CARE)",
+    suffix="_denoised",
+    description="Denoise images using CAREamics (Noise2Void or CARE model)",
+    parameters={
+        "checkpoint_path": {
+            "type": str,
+            "default": "",
+            "description": "Path to the CAREamics model checkpoint file (.ckpt)",
+        },
+        "tile_size_x": {
+            "type": int,
+            "default": 32,
+            "min": 16,
+            "max": 512,
+            "description": "Tile size in X dimension",
+        },
+        "tile_size_y": {
+            "type": int,
+            "default": 32,
+            "min": 16,
+            "max": 512,
+            "description": "Tile size in Y dimension",
+        },
+        "tile_size_z": {
+            "type": int,
+            "default": 0,
+            "min": 0,
+            "max": 256,
+            "description": "Tile size in Z dimension (for 3D data)",
+        },
+        "batch_size": {
+            "type": int,
+            "default": 1,
+            "min": 1,
+            "max": 16,
+            "description": "Batch size for prediction",
+        },
+        "use_tta": {
+            "type": bool,
+            "default": True,
+            "description": "Use test-time augmentation for better results",
+        },
+        "force_dedicated_env": {
+            "type": bool,
+            "default": False,
+            "description": "Force using dedicated environment even if CAREamics is available",
+        },
+    },
+)
+def careamics_denoise(
+    image: np.ndarray,
+    checkpoint_path: str = "",
+    tile_size_z: int = 64,
+    tile_size_y: int = 64,
+    tile_size_x: int = 64,
+    tile_overlap_z: int = 8,
+    tile_overlap_y: int = 8,
+    tile_overlap_x: int = 8,
+    batch_size: int = 1,
+    use_tta: bool = True,
+    force_dedicated_env: bool = False,
+) -> np.ndarray:
+    """
+    Denoise images using CAREamics models.
+
+    This function loads a CAREamics model from a checkpoint file and uses it to denoise
+    the input image. The function supports both 2D and 3D data and handles tiling for
+    processing large images efficiently.
+
+    If CAREamics is not installed in the main environment, a dedicated virtual environment
+    will be automatically created and managed.
+
+    Parameters:
+    -----------
+    image : numpy.ndarray
+        Input image to denoise
+    checkpoint_path : str
+        Path to the CAREamics model checkpoint file (.ckpt)
+    tile_size_z : int
+        Tile size in Z dimension (for 3D data)
+    tile_size_y : int
+        Tile size in Y dimension
+    tile_size_x : int
+        Tile size in X dimension
+    tile_overlap_z : int
+        Tile overlap in Z dimension (for 3D data)
+    tile_overlap_y : int
+        Tile overlap in Y dimension
+    tile_overlap_x : int
+        Tile overlap in X dimension
+    batch_size : int
+        Batch size for prediction
+    use_tta : bool
+        Use test-time augmentation for better results
+    force_dedicated_env : bool
+        Force using dedicated environment even if CAREamics is available
+
+    Returns:
+    --------
+    numpy.ndarray
+        Denoised image with the same dimensions as the input
+    """
+    # Verify checkpoint path
+    if not checkpoint_path or not os.path.exists(checkpoint_path):
+        print(f"Checkpoint file not found: {checkpoint_path}")
+        print("Please provide a valid checkpoint file path.")
+        return image
+
+    # Determine whether to use dedicated environment
+    use_env = force_dedicated_env or USE_DEDICATED_ENV
+
+    careamics_denoise.thread_safe = False
+
+    if use_env:
+        print("Using dedicated CAREamics environment...")
+
+        # First check if the environment exists, create if not
+        if not is_env_created():
+            print(
+                "Creating dedicated CAREamics environment (this may take a few minutes)..."
+            )
+            create_careamics_env()
+            print("Environment created successfully.")
+
+        # Prepare arguments for the CAREamics function
+        args = {
+            "image": image,
+            "checkpoint_path": checkpoint_path,
+            "tile_size_z": tile_size_z,
+            "tile_size_y": tile_size_y,
+            "tile_size_x": tile_size_x,
+            "tile_overlap_z": tile_overlap_z,
+            "tile_overlap_y": tile_overlap_y,
+            "tile_overlap_x": tile_overlap_x,
+            "batch_size": batch_size,
+            "use_tta": use_tta,
+        }
+
+        # Calculate tile overlap automatically (e.g., 25% of tile size)
+        def compute_overlap(tile_size, fraction=0.25):
+            # Ensure overlap is at least 1 and less than tile size
+            overlap = max(1, int(tile_size * fraction))
+            return min(overlap, tile_size - 1)
+
+        # Inside careamics_denoise, after parsing tile sizes:
+        tile_overlap_z = (
+            compute_overlap(tile_size_z) if len(image.shape) >= 3 else None
+        )
+        tile_overlap_y = compute_overlap(tile_size_y)
+        tile_overlap_x = compute_overlap(tile_size_x)
+
+        # Run CAREamics in the dedicated environment
+        print("Running CAREamics in dedicated environment...")
+        return run_careamics_in_env("predict", args)
+
+    else:
+        print("Running CAREamics in current environment...")
+        # Use CAREamics directly in the current environment
+        try:
+            print(f"Loading CAREamics model from: {checkpoint_path}")
+            # Initialize the CAREamist model
+            careamist = CAREamist(
+                checkpoint_path, os.path.dirname(checkpoint_path)
+            )
+
+            # Determine dimensionality
+            is_3d = len(image.shape) >= 3
+
+            if is_3d:
+                print(f"Processing 3D data with shape: {image.shape}")
+                # Determine axes based on dimensionality
+                if len(image.shape) == 3:
+                    # ZYX format
+                    axes = "ZYX"
+                    tile_size = (tile_size_z, tile_size_y, tile_size_x)
+                    tile_overlap = (
+                        tile_overlap_z,
+                        tile_overlap_y,
+                        tile_overlap_x,
+                    )
+                    print(f"Using axes configuration: {axes}")
+                elif len(image.shape) == 4:
+                    # Assuming TZYX format
+                    axes = "TZYX"
+                    tile_size = (tile_size_z, tile_size_y, tile_size_x)
+                    tile_overlap = (
+                        tile_overlap_z,
+                        tile_overlap_y,
+                        tile_overlap_x,
+                    )
+                    print(f"Using axes configuration: {axes}")
+                else:
+                    # Unknown format, try to handle it
+                    print(
+                        f"Warning: Unusual data shape: {image.shape}. Defaulting to 'TZYX'"
+                    )
+                    axes = "TZYX"
+                    tile_size = (tile_size_z, tile_size_y, tile_size_x)
+                    tile_overlap = (
+                        tile_overlap_z,
+                        tile_overlap_y,
+                        tile_overlap_x,
+                    )
+            else:
+                print(f"Processing 2D data with shape: {image.shape}")
+                # 2D data
+                if len(image.shape) == 2:
+                    # YX format
+                    axes = "YX"
+                    tile_size = (tile_size_y, tile_size_x)
+                    tile_overlap = (tile_overlap_y, tile_overlap_x)
+                    print(f"Using axes configuration: {axes}")
+                else:
+                    # Unknown format, try to handle it
+                    print(
+                        f"Warning: Unusual data shape: {image.shape}. Defaulting to 'YX'"
+                    )
+                    axes = "YX"
+                    tile_size = (tile_size_y, tile_size_x)
+                    tile_overlap = (tile_overlap_y, tile_overlap_x)
+
+            # Run the prediction
+            print(
+                f"Running prediction with tile size: {tile_size}, overlap: {tile_overlap}"
+            )
+            print(f"Using batch size: {batch_size}, TTA: {use_tta}")
+
+            prediction = careamist.predict(
+                source=image,
+                tile_size=tile_size,
+                tile_overlap=tile_overlap,
+                axes=axes,
+                batch_size=batch_size,
+                tta=use_tta,
+            )
+
+            # # Handle output shape
+            # if prediction.shape != image.shape:
+            #     print(f"Warning: Prediction shape {prediction.shape} differs from input shape {image.shape}")
+            #     prediction = np.squeeze(prediction)
+
+            #     # If shapes still don't match, try to reshape
+            #     if prediction.shape != image.shape:
+            #         print(f"Warning: Shapes still don't match after squeezing. Using original dimensions.")
+            #         try:
+            #             prediction = prediction.reshape(image.shape)
+            #         except ValueError:
+            #             print("Error: Could not reshape prediction to match input shape.")
+            #             return image
+
+            # print(f"Denoising completed. Output shape: {prediction.shape}")
+            return prediction
+
+        except (RuntimeError, ValueError, ImportError) as e:
+            import traceback
+
+            print(
+                f"Error during CAREamics denoising in current environment: {str(e)}"
+            )
+            traceback.print_exc()
+
+            # If we haven't already tried using the dedicated environment, try that as a fallback
+            if not force_dedicated_env:
+                print(
+                    "Attempting fallback to dedicated CAREamics environment..."
+                )
+                args = {
+                    "image": image,
+                    "checkpoint_path": checkpoint_path,
+                    "tile_size_z": tile_size_z,
+                    "tile_size_y": tile_size_y,
+                    "tile_size_x": tile_size_x,
+                    "tile_overlap_z": tile_overlap_z,
+                    "tile_overlap_y": tile_overlap_y,
+                    "tile_overlap_x": tile_overlap_x,
+                    "batch_size": batch_size,
+                    "use_tta": use_tta,
+                }
+
+                if not is_env_created():
+                    create_careamics_env()
+
+                return run_careamics_in_env("predict", args)
+
+            return None

napari_tmidas/processing_functions/careamics_env_manager.py

@@ -0,0 +1,339 @@
+# processing_functions/careamics_env_manager.py
+"""
+This module manages a dedicated virtual environment for CAREamics.
+"""
+
+import contextlib
+import os
+import platform
+import shutil
+import subprocess
+import tempfile
+import venv
+
+import tifffile
+
+# Define the environment directory in user's home folder
+ENV_DIR = os.path.join(
+    os.path.expanduser("~"), ".napari-tmidas", "envs", "careamics"
+)
+
+
+def is_careamics_installed():
+    """Check if CAREamics is installed in the current environment."""
+    try:
+        import importlib.util
+
+        return importlib.util.find_spec("careamics") is not None
+    except ImportError:
+        return False
+
+
+def is_env_created():
+    """Check if the dedicated environment exists."""
+    env_python = get_env_python_path()
+    return os.path.exists(env_python)
+
+
+def get_env_python_path():
+    """Get the path to the Python executable in the environment."""
+    if platform.system() == "Windows":
+        return os.path.join(ENV_DIR, "Scripts", "python.exe")
+    else:
+        return os.path.join(ENV_DIR, "bin", "python")
+
+
+def create_careamics_env():
+    """Create a dedicated virtual environment for CAREamics."""
+    # Ensure the environment directory exists
+    os.makedirs(os.path.dirname(ENV_DIR), exist_ok=True)
+
+    # Remove existing environment if it exists
+    if os.path.exists(ENV_DIR):
+        shutil.rmtree(ENV_DIR)
+
+    print(f"Creating CAREamics environment at {ENV_DIR}...")
+
+    # Create a new virtual environment
+    venv.create(ENV_DIR, with_pip=True)
+
+    # Path to the Python executable in the new environment
+    env_python = get_env_python_path()
+
+    # Upgrade pip first
+    print("Upgrading pip...")
+    subprocess.check_call(
+        [env_python, "-m", "pip", "install", "--upgrade", "pip"]
+    )
+
+    # Install PyTorch first for compatibility
+    # Try to detect if CUDA is available
+    cuda_available = False
+    try:
+        import torch
+
+        cuda_available = torch.cuda.is_available()
+        print(f"CUDA is {'available' if cuda_available else 'not available'}")
+    except ImportError:
+        print("PyTorch not detected in main environment")
+
+    if cuda_available:
+        # Install PyTorch with CUDA support
+        print("Installing PyTorch with CUDA support...")
+        subprocess.check_call(
+            [env_python, "-m", "pip", "install", "torch", "torchvision"]
+        )
+    else:
+        # Install PyTorch without CUDA
+        print("Installing PyTorch without CUDA support...")
+        subprocess.check_call(
+            [env_python, "-m", "pip", "install", "torch", "torchvision"]
+        )
+
+    # Install CAREamics and dependencies
+    print("Installing CAREamics in the dedicated environment...")
+    subprocess.check_call(
+        [env_python, "-m", "pip", "install", "careamics[tensorboard]"]
+    )
+
+    # Install tifffile for image handling
+    subprocess.check_call(
+        [env_python, "-m", "pip", "install", "tifffile", "numpy"]
+    )
+
+    # Check if installation was successful
+    try:
+        # Run a simple script to verify CAREamics is installed and working
+        check_script = """
+import sys
+try:
+    import careamics
+    print(f"CAREamics version: {careamics.__version__}")
+    from careamics import CAREamist
+    print("CAREamist imported successfully")
+    import torch
+    print(f"PyTorch version: {torch.__version__}")
+    print(f"CUDA available: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"CUDA version: {torch.version.cuda}")
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+    print("SUCCESS: CAREamics environment is working correctly")
+except Exception as e:
+    print(f"ERROR: {str(e)}")
+    sys.exit(1)
+"""
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".py", delete=False
+        ) as temp:
+            temp.write(check_script)
+            temp_path = temp.name
+
+        try:
+            result = subprocess.run(
+                [env_python, temp_path],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            print(result.stdout)
+            if "SUCCESS" in result.stdout:
+                print(
+                    "CAREamics environment created and verified successfully."
+                )
+            else:
+                print(
+                    "WARNING: CAREamics environment created but verification uncertain."
+                )
+        finally:
+            os.unlink(temp_path)
+
+    except subprocess.CalledProcessError as e:
+        print(f"Error during CAREamics installation: {e}")
+        print(f"Output: {e.output}")
+        print(f"Errors: {e.stderr}")
+        raise RuntimeError("Failed to create CAREamics environment") from e
+
+    return env_python
+
+
+def run_careamics_in_env(func_name, args_dict):
+    """
+    Run CAREamics in a dedicated environment.
+
+    Parameters:
+    -----------
+    func_name : str
+        Name of the CAREamics function to run (e.g., 'predict')
+    args_dict : dict
+        Dictionary of arguments for CAREamics function
+
+    Returns:
+    --------
+    numpy.ndarray
+        Denoised image
+    """
+    # Ensure the environment exists
+    if not is_env_created():
+        create_careamics_env()
+
+    # Prepare temporary files
+    with tempfile.NamedTemporaryFile(
+        suffix=".tif", delete=False
+    ) as input_file, tempfile.NamedTemporaryFile(
+        suffix=".tif", delete=False
+    ) as output_file, tempfile.NamedTemporaryFile(
+        mode="w", suffix=".py", delete=False
+    ) as script_file:
+
+        # Save input image to temp file
+        tifffile.imwrite(
+            input_file.name, args_dict["image"], compression="zlib"
+        )
+
+        # Create a temporary script to run CAREamics
+        script = f"""
+import sys
+import os
+import numpy as np
+import tifffile
+from pathlib import Path
+
+try:
+    from careamics import CAREamist
+
+    # Load input image
+    print(f"Loading image from {{os.path.basename('{input_file.name}')}}")
+    image = tifffile.imread('{input_file.name}')
+    print(f"Input image shape: {{image.shape}}, dtype: {{image.dtype}}")
+
+    # Load model
+    print(f"Loading model from {{os.path.basename('{args_dict['checkpoint_path']}')}}")
+    model = CAREamist('{args_dict['checkpoint_path']}',os.path.dirname('{args_dict['checkpoint_path']}'))
+
+    # Determine dimensionality
+    dims = len(image.shape)
+    is_3d = dims >= 3
+
+    if is_3d:
+        print(f"Processing 3D data with shape: {{image.shape}}")
+        # Determine axes based on dimensionality
+        if dims == 3:
+            # ZYX format
+            axes = "ZYX"
+            tile_size = ({args_dict.get('tile_size_z', 64)},
+                         {args_dict.get('tile_size_y', 64)},
+                         {args_dict.get('tile_size_x', 64)})
+            tile_overlap = ({args_dict.get('tile_overlap_z', 32)},
+                           {args_dict.get('tile_overlap_y', 32)},
+                           {args_dict.get('tile_overlap_x', 32)})
+        elif dims == 4:
+            # Assuming TZYX format
+            axes = "TZYX"
+            tile_size = ({args_dict.get('tile_size_z', 64)},
+                         {args_dict.get('tile_size_y', 64)},
+                         {args_dict.get('tile_size_x', 64)})
+            tile_overlap = ({args_dict.get('tile_overlap_z', 32)},
+                           {args_dict.get('tile_overlap_y', 32)},
+                           {args_dict.get('tile_overlap_x', 32)})
+        else:
+            print(f"Warning: Unusual data shape: {{image.shape}}. Defaulting to 'TZYX'")
+            axes = "TZYX"
+            tile_size = ({args_dict.get('tile_size_z', 64)},
+                         {args_dict.get('tile_size_y', 64)},
+                         {args_dict.get('tile_size_x', 64)})
+            tile_overlap = ({args_dict.get('tile_overlap_z', 32)},
+                           {args_dict.get('tile_overlap_y', 32)},
+                           {args_dict.get('tile_overlap_x', 32)})
+    else:
+        print(f"Processing 2D data with shape: {{image.shape}}")
+        # 2D data
+        if dims == 2:
+            # YX format
+            axes = "YX"
+            tile_size = ({args_dict.get('tile_size_y', 64)},
+                         {args_dict.get('tile_size_x', 64)})
+            tile_overlap = ({args_dict.get('tile_overlap_y', 32)},
+                           {args_dict.get('tile_overlap_x', 32)})
+        else:
+            print(f"Warning: Unusual data shape: {{image.shape}}. Defaulting to 'YX'")
+            axes = "YX"
+            tile_size = ({args_dict.get('tile_size_y', 64)},
+                         {args_dict.get('tile_size_x', 64)})
+            tile_overlap = ({args_dict.get('tile_overlap_y', 32)},
+                           {args_dict.get('tile_overlap_x', 32)})
+
+    print(f"Using axes: {{axes}}, tile_size: {{tile_size}}, tile_overlap: {{tile_overlap}}")
+
+    # Run the prediction
+    print("Starting CAREamics prediction...")
+    prediction = model.predict(
+        source=image,
+        tile_size=tile_size,
+        tile_overlap=tile_overlap,
+        axes=axes,
+        batch_size={args_dict.get('batch_size', 1)},
+        tta={args_dict.get('use_tta', True)},
+    )
+
+    # # Handle output shape
+    # if prediction.shape != image.shape:
+    #     print(f"Warning: Prediction shape {{prediction.shape}} differs from input shape {{image.shape}}")
+    #     prediction = np.squeeze(prediction)
+
+    #     # If shapes still don't match, try to reshape
+    #     if prediction.shape != image.shape:
+    #         print(f"Warning: Shapes still don't match after squeezing. Using original dimensions.")
+
+    # print(f"Denoising completed. Output shape: {{prediction.shape}}")
+
+    # Save result
+    print(f"Saving result to {{os.path.basename('{output_file.name}')}}")
+    tifffile.imwrite('{output_file.name}', prediction, compression="zlib")
+    print("Done!")
+
+    # Exit with success code
+    sys.exit(0)
+
+except Exception as e:
+    import traceback
+    print(f"Error in CAREamics processing: {{e}}")
+    traceback.print_exc()
+    sys.exit(1)
+"""
+
+        # Write script
+        script_file.write(script)
+        script_file.flush()
+
+    try:
+        # Run the script in the dedicated environment
+        env_python = get_env_python_path()
+        result = subprocess.run(
+            [env_python, script_file.name], capture_output=True, text=True
+        )
+
+        # Print output
+        print(result.stdout)
+
+        # Check for errors
+        if result.returncode != 0:
+            print("Error in CAREamics processing:")
+            print(result.stderr)
+            raise RuntimeError(
+                f"CAREamics denoising failed with error code {result.returncode}"
+            )
+
+        # Read and return the results
+        denoised_image = tifffile.imread(output_file.name)
+        return denoised_image
+
+    except (RuntimeError, subprocess.CalledProcessError) as e:
+        print(f"Error in CAREamics processing: {e}")
+        # Return original image in case of error
+        return args_dict["image"]
+
+    finally:
+        # Clean up temporary files
+        for file_path in [input_file.name, output_file.name, script_file.name]:
+            with contextlib.suppress(FileNotFoundError, PermissionError):
+                os.unlink(file_path)