inference.nnunet.install_model

Import Model from ZIP File

Provides functionalities related to extracting a pretrained nnU-Net model from a zipped file.

View Source

 1"""
 2### Import Model from ZIP File
 3
 4Provides functionalities related to extracting a pretrained nnU-Net model from a zipped file.
 5"""
 6
 7
 8import zipfile
 9import os
10import argparse
11from typing import Optional
12
13
14def check_path(path: str) -> None:
15    """
16    Checks if the provided path points to a ZIP file.
17
18    Args:
19        path (str): Path to be checked.
20
21    Raises:
22        AssertionError: If the path does not have a ".zip" extension.
23    """
24    assert path.endswith(".zip"), "Path must point to a ZIP file"
25
26
27def install_model_from_zip(path: str, mednext: bool = False) -> None:
28    """
29    Installs a pretrained nnU-Net (v2) model from a zipped file.
30
31    Args:
32        path (str): Path to the zipped model file.
33        mednext (bool, optional): Flag to specify if the MedNext configuration is used. Defaults to False.
34
35    Raises:
36        AssertionError: If the provided path is not a ZIP file.
37    """
38    # Verify that the provided path is a ZIP file
39    check_path(path)
40    
41    # Default extraction path
42    zip_path = '/tmp/'
43
44    if mednext:
45        # If MedNext is specified, change the extraction path and create the directory if it doesn't exist
46        zip_path = '/tmp/nnUNet'
47        os.makedirs(zip_path, exist_ok=True)
48    
49    # Extract the ZIP file to the specified directory
50    with zipfile.ZipFile(path, 'r') as zip_ref:
51        zip_ref.extractall(zip_path)
52    
53    # Uncomment the following lines if you need to run an external command after extraction
54    # cmd = f"nnUNetv2_install_pretrained_model_from_zip {path}"
55    # subprocess.run(cmd, shell=True)  # Executes the command in the shell

def check_path(path: str) -> None: View Source

15def check_path(path: str) -> None:
16    """
17    Checks if the provided path points to a ZIP file.
18
19    Args:
20        path (str): Path to be checked.
21
22    Raises:
23        AssertionError: If the path does not have a ".zip" extension.
24    """
25    assert path.endswith(".zip"), "Path must point to a ZIP file"

Checks if the provided path points to a ZIP file.

Args: path (str): Path to be checked.

Raises: AssertionError: If the path does not have a ".zip" extension.

def install_model_from_zip(path: str, mednext: bool = False) -> None: View Source

28def install_model_from_zip(path: str, mednext: bool = False) -> None:
29    """
30    Installs a pretrained nnU-Net (v2) model from a zipped file.
31
32    Args:
33        path (str): Path to the zipped model file.
34        mednext (bool, optional): Flag to specify if the MedNext configuration is used. Defaults to False.
35
36    Raises:
37        AssertionError: If the provided path is not a ZIP file.
38    """
39    # Verify that the provided path is a ZIP file
40    check_path(path)
41    
42    # Default extraction path
43    zip_path = '/tmp/'
44
45    if mednext:
46        # If MedNext is specified, change the extraction path and create the directory if it doesn't exist
47        zip_path = '/tmp/nnUNet'
48        os.makedirs(zip_path, exist_ok=True)
49    
50    # Extract the ZIP file to the specified directory
51    with zipfile.ZipFile(path, 'r') as zip_ref:
52        zip_ref.extractall(zip_path)
53    
54    # Uncomment the following lines if you need to run an external command after extraction
55    # cmd = f"nnUNetv2_install_pretrained_model_from_zip {path}"
56    # subprocess.run(cmd, shell=True)  # Executes the command in the shell

Installs a pretrained nnU-Net (v2) model from a zipped file.

Args: path (str): Path to the zipped model file. mednext (bool, optional): Flag to specify if the MedNext configuration is used. Defaults to False.

Raises: AssertionError: If the provided path is not a ZIP file.