Brain tumor image automatic segmentation method, system, terminal and medium

Abstract

The present invention provides a method and system for automatic segmentation of brain tumor images, wherein the method includes: constructing a first brain tumor segmentation model based on Pytorch coding and nnUNet framework; converting the first brain tumor segmentation model from Pytorch format to ONNX format to obtain a second brain tumor segmentation model; visualizing the second brain tumor segmentation model, and pruning the second brain tumor segmentation model based on the visualized model structure to obtain a third brain tumor segmentation model; quantizing the third brain tumor segmentation model through the TensorRT optimization tool to obtain a brain tumor image automatic segmentation model for automatic brain tumor image segmentation. The present invention greatly improves the inference speed of image segmentation through the TensorRT model without affecting the basic performance of the model, reduces the inference time, and improves work efficiency.

Description

Technical Field

[0001] The present invention relates to the technical field of medical image segmentation, and in particular to a method, system, terminal and medium for automatic segmentation of brain tumor images based on TensorRT accelerated reasoning. Background Art

[0002] Medical image segmentation technology is a complex and critical step in medical image processing and analysis. Its purpose is to segment and process the parts with special meanings (lesion parts) in medical images and extract relevant features to provide a reliable basis for clinical diagnosis and pathological research, and assist doctors in making more accurate judgments.

[0003] The existing technology for segmenting medical images can be directly implemented through model inference using a deep learning framework. However, the existing inference technology has limited speed, especially for brain MRI images with more details and more complex structures. It requires longer inference time and is unable to prune and quantize the segmentation model. It has strict requirements on the version of the dependent library and the system, and cannot meet the increasing technical demand for rapid segmentation of brain tumors in brain medical images. Summary of the invention

[0004] In view of the above-mentioned deficiencies in the prior art, the present invention provides a method, system, terminal and medium for automatic segmentation of brain tumor images.

[0005] According to one aspect of the present invention, a method for automatic segmentation of brain tumor images is provided, comprising:

[0006] Build the first brain tumor segmentation model based on Pytorch coding and nnUNet framework;

[0007] Converting the first brain tumor segmentation model from a Pytorch format to an ONNX format to obtain a second brain tumor segmentation model;

[0008] Visualizing the second brain tumor segmentation model, and based on the visualized model structure, pruning the second brain tumor segmentation model to obtain a third brain tumor segmentation model;

[0009] The third brain tumor segmentation model is quantized by the TensorRT optimization tool to obtain a brain tumor image automatic segmentation model for automatic brain tumor image segmentation.

[0010] Preferably, the first brain tumor segmentation model is constructed based on Pytorch coding and nnUNet framework, including:

[0011] Using Pytorch coding, image preprocessing and image enhancement operations were performed on brain MRI images and corresponding images labeled with brain tumor labels to obtain 3D image data;

[0012] Performing three-dimensional sliding window clipping on the 3D image data to obtain training samples;

[0013] Providing a nnUNet network model, and training the nnUNet network model using the training samples;

[0014] Performing a preliminary pre-training loss analysis in the nnUNet network model according to the training samples, selecting a corresponding network depth and number of convolution kernels for losses in different ranges and downward trends, and obtaining an adaptive network parameter configuration;

[0015] By configuring the adaptive network parameters, relatively feasible network hyperparameters are obtained, and the nnUNet network model is optimized according to the network hyperparameters to obtain a brain tumor segmentation network model;

[0016] The brain MRI image is used as the input of the brain tumor segmentation network model, the brain tumor segmentation network model is used for inference, and the inference result is subjected to small noise processing to obtain a predicted brain tumor label, and the brain tumor segmentation network model is tested to obtain the first brain tumor segmentation model in Pytorch format.

[0017] Preferably, converting the first brain tumor segmentation model from Pytorch format to ONNX format to obtain a second brain tumor segmentation model includes:

[0018] Acquire the network structure and network parameters of the first brain tumor segmentation model;

[0019] The network structure and the network parameters are passed into the torch.onnx.export model to obtain a second brain tumor segmentation model in ONNX format;

[0020] The dynamic_axes parameter is added to the second brain tumor segmentation model in the ONNX format to set the input and output to dynamic dimensions so that the shape and size of the input and output can change in different inference requests.

[0021] Preferably, the visualizing the second brain tumor segmentation model and pruning the second brain tumor segmentation model based on the visualized model structure to obtain a third brain tumor segmentation model includes:

[0022] Using the Netron visualization tool, generating a visualization model structure for the second brain tumor segmentation model;

[0023] Confirming useless outputs in the visualization model structure, and obtaining node names, weight names and / or serial numbers corresponding to the useless outputs;

[0024] Traversing and deleting the useless outputs in the second brain tumor segmentation model by using the node names, weight names and / or serial numbers;

[0025] The deleted model is saved to obtain the third brain tumor segmentation model.

[0026] Preferably, the third brain tumor segmentation model is quantized by a TensorRT optimization tool to obtain an automatic brain tumor image segmentation model, including:

[0027] The third brain tumor segmentation model in ONNX format is converted into a TensorRT format model by using the trtexec program of the TensorRT optimizer, and quantization parameters are selected for the TensorRT format model to obtain a brain tumor image automatic segmentation model in TensorRT format.

[0028] Preferably, the quantization parameter selection includes: selecting fp16 or int8 as the quantization parameter.

[0029] Preferably, the method uses the trtexec program of the TensorRT optimizer to convert the third brain tumor segmentation model in ONNX format into a TensorRT format model, and selects quantization parameters for the TensorRT format model to obtain a brain tumor image automatic segmentation model in TensorRT format, and further includes:

[0030] Use the trtexec program of the TensorRT optimizer to build a TensorRT reasoning class, which is used to encapsulate the TensorRT model segmentation reasoning process;

[0031] Deserialize a TensorRT engine from the given ONNX format model file path of the third brain tumor segmentation model, create the TensorRT engine execution context, and allocate host memory and device memory for input and output of the TensorRT engine;

[0032] Define a memory allocation method and an inference execution method; wherein: the memory allocation method is used to allocate input and output host memory and device memory and related buffers for the TensorRT engine, and traverse all bindings of the TensorRT engine to allocate corresponding host memory and device memory according to the input or output type of the binding, and add the device memory to the binding list; the inference execution method is used to execute the segmentation inference process, wherein the segmentation inference process includes: first setting the input data to the passed array and executing the TensorRT inference, then extracting the output image from the inference result and returning the inference result;

[0033] Instantiate the TensorRT inference class, read in the quantized TensorRT model, and build a brain tumor image automatic segmentation model in TensorRT format;

[0034] By using the brain tumor image automatic segmentation model in the TensorRT format, the brain MRI image to be segmented is input into the TensorRT reasoning class of the brain tumor image automatic segmentation model, and the segmentation reasoning process is executed to perform automatic segmentation of the brain tumor image.

[0035] Preferably, the above method further comprises:

[0036] Build a Docker container environment, which is used to provide a model building environment.

[0037] Preferably, the above method further comprises:

[0038] The Flask framework is used to encapsulate the brain tumor image automatic segmentation model.

[0039] Preferably, the use of the Flask framework to encapsulate the brain tumor image automatic segmentation model includes:

[0040] The brain tumor image automatic segmentation model is encapsulated by using a Flask framework to obtain a Flask application;

[0041] Use Dockerfile to build an image, and start the container with the built image to expose the port number and increase the memory limit at startup;

[0042] Open the GPU and automatically start the Flask application when the image startup container is started, which is used to externally call the corresponding Flask brain tumor image automatic segmentation service.

[0043] According to another aspect of the present invention, there is provided a brain tumor image automatic segmentation system, comprising:

[0044] The Pytorch model module builds the first brain tumor segmentation model based on Pytorch coding and the nnUNet framework;

[0045] An ONNX model module, the module being used to convert the first brain tumor segmentation model from a Pytorch format to an ONNX format to obtain a second brain tumor segmentation model;

[0046] a pruning module, the module being used to visualize the second brain tumor segmentation model, and based on the visualized model structure, prune the second brain tumor segmentation model to obtain a third brain tumor segmentation model;

[0047] A TensorRT model module is used to quantize the third brain tumor segmentation model through the TensorRT optimization tool to obtain an automatic brain tumor image segmentation model for automatic brain tumor image segmentation.

[0048] Preferably, the above system further includes:

[0049] An environment building module, which is used to build a Docker container environment, and the Docker container environment is used to build the Pytorch model module, the ONNX model module, the pruning module and the TensorRT model module.

[0050] Preferably, the above system further includes:

[0051] The Flask application module uses the Flask framework to encapsulate the brain tumor image automatic segmentation model.

[0052] According to a third aspect of the present invention, there is provided a computer terminal comprising a memory, a processor and a computer program stored in the memory and executable on the processor, wherein the processor, when executing the computer program, can be used to execute any of the above-described methods of the present invention, or to run any of the above-described systems of the present invention.

[0053] According to a fourth aspect of the present invention, there is provided a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, can be used to execute any of the methods described above in the present invention, or to run any of the systems described above in the present invention.

[0054] Due to the adoption of the above technical solution, the present invention has at least one of the following beneficial effects compared with the prior art:

[0055] The present invention provides a method, system, terminal and medium for automatic brain tumor image segmentation. A brain tumor segmentation model trained based on the Pytorch and nnUNet deep learning framework is converted into an ONNX model through the onnx.export function to achieve model visualization and pruning. The ONNX model is then quantized and converted through the TensorRT optimization tool. The size of the final TensorRT model is reduced. Without affecting the basic performance of the model, the inference speed of image segmentation completed by the TensorRT model is greatly improved, the inference time is reduced, and the work efficiency is improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0056] Other features, objects and advantages of the present invention will become more apparent from the detailed description of non-limiting embodiments made with reference to the following drawings:

[0057] Figure 1 FIG. 4 is a flowchart of a method for automatic segmentation of brain tumor images in one embodiment of the present invention.

[0058] Figure 2 Schematic diagram of the overall working process of the brain tumor image automatic segmentation method in a preferred embodiment of the present invention.

[0059] Figure 3 Schematic diagram of model training for a method for automatic segmentation of brain tumor images in a preferred embodiment of the present invention.

[0060] Figure 4 Schematic diagram of the components of the brain tumor image automatic segmentation system in one embodiment of the present invention. DETAILED DESCRIPTION

[0061] The following is a detailed description of the embodiments of the present invention: This embodiment is implemented on the premise of the technical solution of the present invention, and a detailed implementation method and a specific operation process are given. It should be pointed out that for ordinary technicians in this field, several variations and improvements can be made without departing from the concept of the present invention, which all belong to the protection scope of the present invention.

[0062] Brain tumor segmentation based on brain MRI images can provide a reliable basis for clinical diagnosis and pathological research, and is an important technical means to assist doctors in making more accurate judgments. Existing image segmentation technology, especially in brain tumor segmentation based on brain MRI images, usually has the problems of long reasoning time and low work efficiency.

[0063] In response to the above problems, an embodiment of the present invention provides a method for automatic segmentation of brain tumor images. The method converts a brain tumor segmentation model trained based on the Pytorch and nnUNet deep learning frameworks into an ONNX model through the onnx.export function. At this time, the model can be visualized and pruned, and then quantized and converted through TensorRT. The obtained TensorRT model can realize automatic segmentation of brain tumors with accelerated reasoning.

[0064] Specifically, Figure 1 As shown, the brain tumor image automatic segmentation method provided in this embodiment may include the following operations:

[0065] S1, build the first brain tumor segmentation model based on Pytorch coding and nnUNet framework;

[0066] S2, converting the first brain tumor segmentation model from Pytorch format to ONNX format to obtain a second brain tumor segmentation model;

[0067] S3, visualizing the second brain tumor segmentation model, and based on the visualized model structure, pruning the second brain tumor segmentation model to obtain a third brain tumor segmentation model;

[0068] S4, quantizing the third brain tumor segmentation model through the TensorRT optimization tool to obtain a brain tumor image automatic segmentation model for automatic brain tumor image segmentation.

[0069] In some preferred embodiments, the above method may further include the following operations:

[0070] S0: Build a Docker container environment. The Docker container environment is used to execute S1 to S4 above and provide an environment for building the model.

[0071] In order to obtain a brain tumor segmentation model and provide a basis for ONNX format conversion, in some preferred implementations, the above S1, based on Pytorch encoding and nnUNet framework, constructs a first brain tumor segmentation model, and may further include the following operations:

[0072] S11, using Pytorch coding, image preprocessing and image enhancement operations were performed on brain MRI images and corresponding images annotated with brain tumor labels to obtain 3D image data;

[0073] S12, performing three-dimensional sliding window clipping on the 3D image data to obtain training samples;

[0074] S13, providing a nnUNet network model, and training the nnUNet network model using the training samples;

[0075] S14, in the nnUNet network model, a preliminary pre-training loss analysis is performed based on the training samples, and the corresponding network depth and number of convolution kernels are selected according to the loss in different ranges and the downward trend, so as to obtain an adaptive network parameter configuration; by calculating the degree and speed of decline of the training loss, different ranges correspond to different depths and numbers of convolution kernels, so that the appropriate UNet network depth and number of convolution kernels in the network structure can be found, and finally the appropriate network parameters can be obtained;

[0076] S15, obtaining relatively feasible network hyperparameters through adaptive network parameter configuration, and optimizing the nnUNet network model according to the network hyperparameters to obtain a brain tumor segmentation network model;

[0077] S16, using the brain MRI image as the input of the brain tumor segmentation network model, inferring through the brain tumor segmentation network model, and performing small noise processing on the inference results to obtain the predicted brain tumor label, testing the brain tumor segmentation network model, and obtaining the first brain tumor segmentation model in Pytorch format. In this step, the small noise is a false positive predicted by the network or a small set of pixel points predicted by the network; the post-processing method of removing small noise from the inference results can obtain a clearer image result.

[0078] In some preferred implementations, the above S11, image preprocessing, may further include the following operations:

[0079] The brain MRI images and the corresponding images annotated with brain tumor labels were corrected, standardized and normalized using the N4 bias field.

[0080] In some preferred implementations, the above S11, image enhancement operation, may further include the following operations:

[0081] The preprocessed brain MRI images and the corresponding images annotated with brain tumor labels are rotated, flipped, scaled, translated and / or contrast transformed.

[0082] In some preferred implementations, the above S14, selecting corresponding network depths and numbers of convolution kernels for losses within different ranges and downward trends to obtain adaptive network parameter configurations, may further include the following operations:

[0083] S141, first give an initial network parameter, including: learning rate, batch size, number of network layers, number of convolution kernels;

[0084] S142, then automatically calculate the average slope based on the training epochs and the corresponding dice loss. If 0 > slope > a, it means the decrease is too slow, and the learning rate needs to be increased accordingly. If the slope < a, it means the decrease is too fast, and the learning rate needs to be decreased accordingly. Here, the adjustment of the learning rate can be adjusted according to the slope value. In a specific application example of this step, the parameter a can take the value of -0.5;

[0085] S143, the number of network layers and the number of convolutional kernels are judged by the decrease value of the dice loss in the first n epochs. If the dice loss drops below b, the network is prone to overfitting. At this time, the network depth is reduced by one layer, and then trained for n more epochs. Then, it is judged again whether the dice loss drops below b. If it drops below b, the network depth is reduced by one more layer, and the network depth is reduced by at most two layers in total. In a specific application example of this step, the parameter n can take the value of 5, and / or the parameter b can take the value of -0.3;

[0086] S144, the Batch size is automatically calculated according to the GPU video memory and occupancy, and the maximum value is 8, so that the video memory occupancy rate is between 50 - 70%, so as to automatically find a suitable batch size.

[0087] In some preferred embodiments, for the above S16, and perform small noise processing on the inference result, the connected components smaller than a certain size (set value) can be removed by the morphology.remove_small_objects method of the skimage library.

[0088] In order to obtain a general and easily optimized model, ONNX is selected as the intermediate process model to prepare for the quantization and inference of the TenorRT model. In some preferred embodiments, for the above S2, convert the first brain tumor segmentation model from the Pytorch format to the ONNX format to obtain the second brain tumor segmentation model. It may further include the following operations:

[0089] S21, obtain the network structure and network parameters of the first brain tumor segmentation model;

[0090] S22, pass the network structure and network parameters into the torch.onnx.export model to obtain the second brain tumor segmentation model in the ONNX format;

[0091] Among them, the dynamic_axes parameter is added to the second brain tumor segmentation model in the ONNX format to set the input and output as dynamic dimensions, so that the shape and size of the input and output can vary in different inference requests for automatically segmenting brain tumors in images.

[0092] In order to clearly prune the brain tumor segmentation model, in some preferred embodiments, the above S3, visualizing the second brain tumor segmentation model, and pruning the second brain tumor segmentation model based on the visualized model structure to obtain the third brain tumor segmentation model, may further include the following operations:

[0093] S31, using the Netron visualization tool, generate a visualization model structure for the second brain tumor segmentation model;

[0094] S32, confirming useless outputs in the visualization model structure, and obtaining node names, weight names and / or serial numbers corresponding to the useless outputs; wherein different nodes of the model may have many outputs, and except for the required final output, other outputs may be regarded as useless outputs;

[0095] S33, traversing and deleting useless outputs in the second brain tumor segmentation model by node name, weight name and / or sequence number;

[0096] S34, saving the deleted model to obtain a third brain tumor segmentation model.

[0097] In order to quantify the brain tumor segmentation model, in some preferred embodiments, the above S4, quantizing the third brain tumor segmentation model through the TensorRT optimization tool to obtain the brain tumor image automatic segmentation model, may further include the following operations:

[0098] The third brain tumor segmentation model in ONNX format was converted into a TensorRT format model using the trtexec program of the TensorRT optimizer. Then, quantization parameters were selected for the TensorRT format model to obtain an automatic brain tumor image segmentation model in TensorRT format.

[0099] In some preferred implementations, the above quantization parameter selection may further include the following operations:

[0100] Select fp16 or int8 as the quantization parameter.

[0101] In some preferred implementations, the above S4 may further include the following operations:

[0102] S41, use the trtexec program of the TensorRT optimizer to build the TensorRT reasoning class, which is used to encapsulate the TensorRT model segmentation reasoning process;

[0103] S42, deserialize a TensorRT engine from a given model file path of the third brain tumor segmentation model in ONNX format, create a TensorRT engine execution context, and allocate host memory and device memory for input and output of the TensorRT engine;

[0104] S43, defining a memory allocation method and an inference execution method; wherein: the memory allocation method is used to allocate input and output host memory and device memory and related buffers for the TensorRT engine, by traversing all bindings of the TensorRT engine, allocating corresponding host memory and device memory according to the input or output type of the binding, and adding the device memory to the binding list; the inference execution method is used to execute the segmentation inference process, wherein the segmentation inference process includes: first setting the input data to the passed-in array and executing TensorRT inference, then extracting the output image from the inference result, adjusting it according to the special situation of the model, and finally returning the inference result;

[0105] S44, instantiate the TensorRT reasoning class, read in the quantized TensorRT model, and construct a brain tumor image automatic segmentation model in the TensorRT format;

[0106] S45, using the brain tumor image automatic segmentation model in TensorRT format, inputting the brain MRI image to be segmented into the TensorRT reasoning class of the brain tumor image automatic segmentation model, executing the segmentation reasoning process, and performing automatic segmentation of the brain tumor image.

[0107] In some preferred implementations, in the above S42, a TensorRT engine is deserialized from the model file path of the third brain tumor segmentation model in a given ONNX format, and the deserialize_cuda_engine() method in the TensorRT library can be used to deserialize the TensorRT engine.

[0108] In order to provide an external call service and automatically start the service, in some preferred implementations, the above method may further include the following operations:

[0109] S5, the Flask framework is used to encapsulate the brain tumor image automatic segmentation model.

[0110] In some preferred implementations, the above S5, using the Flask framework to encapsulate the brain tumor image automatic segmentation model, may further include the following operations:

[0111] S51, encapsulating the brain tumor image automatic segmentation model by using the Flask framework to obtain a Flask application;

[0112] S52, uses Dockerfile to build an image, and starts the container through the built image to expose the port number and increase the memory limit at startup;

[0113] S53, opens the GPU and automatically starts the Flask application when the image startup container is started, which is used to externally call the corresponding Flask brain tumor image automatic segmentation service.

[0114] The technical solution provided by the above embodiment of the present invention is further described in detail below in conjunction with a specific application example.

[0115] In the specific application example, taking the automatic segmentation of meningioma in brain MRI images as an example, the following steps are included:

[0116] like Figure 2 As shown, first, a Docker container is built to manage the entire dependent environment, and the Pytorch encoding and nnUNet framework are used to train the first brain tumor segmentation model in Pytorch format. The Pytorch format model is then converted into an ONNX format model. The overall model structure can be visualized through Netron to find unnecessary or insignificant model branches for pruning. The ONNX format model is then quantized through the TensorRT optimization tool and converted into a TensorRT format model. The TensorRT format model is used for automatic segmentation of meningioma images. Finally, the model can be encapsulated in a Flask service for external calls to the TensorRT format brain tumor image automatic segmentation model in the Flask service for automatic segmentation of meningioma.

[0117] Further:

[0118] Docker container environment construction:

[0119] Install the appropriate version of CUDA according to the GPU model of the host, pull the corresponding CUDA version image from the official Docker, install the appropriate versions of TensorRT, Pycuda, ONNX, Flask, Pytorch, nnUNet and other related medical image processing libraries. The preprocessing, training and inference modules all need to be run in this environment.

[0120] Brain tumor segmentation model training based on Pytorch and nnUNet:

[0121] The brain MRI images and the corresponding images labeled with meningioma labels are preprocessed and trained using the nnUNet framework. The specific parameters and configurations can be optimized through multiple experiments, and the trained model is saved in Pytorch format.

[0122] The training process is as follows Figure 3 As shown, firstly, the prepared brain MRI images and the corresponding images labeled with meningioma labels are preprocessed, and the preprocessing includes: N4 bias field correction, standardization, and normalization. The preprocessed brain MRI images and the corresponding images labeled with meningioma labels are then subjected to image enhancement operations, including: rotation, flipping, scaling, translation, contrast transformation, etc. Since the acquired data is a 3D image, the enhanced brain MRI images and the corresponding images labeled with meningioma labels can be subjected to three-dimensional sliding window cropping. The cropped brain MRI images are used as the input of the nnUNet network model, and the corresponding images labeled with meningioma labels are The image is used as the output of the nnUNet network model to train the nnUNet network model; in this network module, a preliminary pre-training loss analysis is further performed based on the image, and then the corresponding network depth and number of convolution kernels are selected for the loss in different ranges and the downward trend to obtain an adaptive network parameter configuration. The network depth in the network parameters can be detected by detecting the network downsampling depth and number of times through the python code. At the same time, relatively feasible network hyperparameters are obtained through adaptive network parameter configuration. The nnUNet network model is further optimized according to the appropriate network hyperparameters to obtain a brain tumor segmentation network model. The brain MRI image is used as the input of the brain tumor segmentation network model, and the meningioma in the image is predicted and small noise is removed, and the predicted meningioma label is output.

[0123] Convert the Pytorch model to an ONNX model:

[0124] Get the network structure and parameters of the trained brain tumor segmentation model in Pytorch format, and export the second brain tumor segmentation model in ONNX format by passing the network structure and network parameters into the torch.onnx.export model. The dynamic_axes parameter is added to the model to set the input and output as dynamic dimensions so that their sizes can change in different inference requests.

[0125] Visualize ONNX model and prune:

[0126] The ONNX model has some redundant and useless outputs. You can use Netron to visualize the overall model structure, find the node names, weight names, and serial numbers of these useless outputs, and then import the ONNX model, traverse these names and serial numbers, delete the redundant parts, and save the ONNX model again.

[0127] Convert ONNX model to TensorRT quantized model:

[0128] The ONNX model can be directly converted to the TensorRT model through the trtexec program in the installed TensorRT folder. The quantization parameter can be selected as fp16 or int8, which can also reduce the model size.

[0129] TensorRT Model:

[0130] Build the TensorRT inference class, which encapsulates the TensorRT model segmentation inference process. In the constructor, the class deserializes a TensorRT engine from the given model file path, creates the TensorRT engine execution context, and allocates the host memory and device memory for the input and output of the TensorRT engine. Define the memory allocation method and the execution inference method. Among them: the memory allocation method is used to allocate the host memory and device memory and related buffers for the input and output of the TensorRT engine. It traverses all bindings of the TensorRT engine, allocates the corresponding host memory and device memory according to the type of binding (input or output), and adds the device memory to the binding list. The execution inference method is used to execute the segmentation inference process, which includes: first setting the input data to the passed array and performing TensorRT inference, then extracting the output data from the inference result, adjusting it according to the special situation of the model, and finally returning the inference result. Finally, the TensorRT reasoning class is instantiated, the TensorRT quantization model is read in, and the brain tumor image automatic segmentation model in TensorRT format is constructed. The brain MRI image to be segmented for meningioma is read and put into the TensorRT reasoning class, the segmentation reasoning process is executed, the meningioma image is automatically segmented, and the reasoning result is returned.

[0131] Flask service:

[0132] Configure the environment, model, and inference code using the above method. You can also install anaconda to manage the environment for different Python versions, use Dockerfile to build an image, and then start the container using the image built above. When starting, set the port number exposed to the outside and increase the memory limit. At the same time, open the GPU. When the container starts, the Flask app will be automatically started for external calls to the Flask meningioma image automatic segmentation service of the container.

[0133] Through application verification in specific application examples, the brain tumor image automatic segmentation method provided by the above embodiment of the present invention is used to perform meningioma segmentation. Compared with the existing medical image segmentation method, the inference speed is increased by more than half.

[0134] An embodiment of the present invention provides a brain tumor image automatic segmentation system. Figure 4As shown, the system may include the following modules:

[0135] The Pytorch model module builds the first brain tumor segmentation model based on Pytorch coding and the nnUNet framework;

[0136] An ONNX model module, the module is used to convert the first brain tumor segmentation model from a Pytorch format to an ONNX format to obtain a second brain tumor segmentation model;

[0137] A pruning module, which is used to visualize the second brain tumor segmentation model and prune the second brain tumor segmentation model based on the visualized model structure to obtain a third brain tumor segmentation model;

[0138] The TensorRT model module is used to quantize the third brain tumor segmentation model through the TensorRT optimization tool to obtain an automatic brain tumor image segmentation model for automatic brain tumor image segmentation.

[0139] In some preferred embodiments, the above system may further include any one or more of the following modules:

[0140] -Environment building module, which is used to build the Docker container environment, which is used to build the Pytorch model module, ONNX model module, pruning module and TensorRT model module;

[0141] -Flask application module, which uses the Flask framework to encapsulate the brain tumor image automatic segmentation model.

[0142] It should be noted that the steps in the method provided by the present invention can be implemented by using corresponding modules, devices, units, etc. in the system. Those skilled in the art can refer to the technical solution of the method to realize the composition of the system, that is, the embodiments in the method can be understood as preferred examples for constructing the system, which will not be elaborated here.

[0143] An embodiment of the present invention provides a computer terminal, including a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, it can be used to execute any method of the above-mentioned embodiments of the present invention, or to run any system of the above-mentioned embodiments of the present invention.

[0144] Optionally, the memory is used to store programs; the memory may include volatile memory (English: volatile memory), such as random-access memory (English: random-access memory, abbreviated: RAM), such as static random-access memory (English: static random-access memory, abbreviated: SRAM), double data rate synchronous dynamic random access memory (English: Double Data Rate Synchronous Dynamic Random Access Memory, abbreviated: DDR SDRAM), etc.; the memory may also include non-volatile memory (English: non-volatile memory), such as flash memory (English: flash memory). The memory is used to store computer programs (such as applications, functional modules, etc. that implement the above method), computer instructions, etc., and the above computer programs, computer instructions, etc. can be partitioned and stored in one or more memories. And the above computer programs, computer instructions, data, etc. can be called by the processor.

[0145] The above-mentioned computer programs, computer instructions, etc. may be stored in one or more memories in partitions, and the above-mentioned computer programs, computer instructions, data, etc. may be called by a processor.

[0146] The processor is used to execute the computer program stored in the memory to implement the various steps of the method or various modules of the system involved in the above embodiments. For details, please refer to the relevant descriptions in the above method and system embodiments.

[0147] The processor and the memory may be independent structures or integrated structures. When the processor and the memory are independent structures, the memory and the processor may be coupled and connected via a bus.

[0148] An embodiment of the present invention provides a computer-readable storage medium having a computer program stored thereon. When the computer program is executed by a processor, it can be used to execute any method of the above embodiments of the present invention, or to run any system of the above embodiments of the present invention.

[0149] Those skilled in the art know that, in addition to implementing the system and its various devices provided by the present invention in a purely computer-readable program code, the system and its various devices provided by the present invention can be made to implement the same functions in the form of logic gates, switches, application-specific integrated circuits, programmable logic controllers, and embedded microcontrollers by logically programming the method steps. Therefore, the system and its various devices provided by the present invention can be considered as a hardware component, and the devices for implementing various functions included therein can also be considered as structures within the hardware component; the devices for implementing various functions can also be considered as both software modules for implementing the method and structures within the hardware component.

[0150] All matters not covered in the above embodiments of the present invention are well known in the art.

[0151] The above describes the specific embodiments of the present invention. It should be understood that the present invention is not limited to the above specific embodiments, and those skilled in the art may make various modifications or variations within the scope of the claims, which do not affect the essence of the present invention.

Claims

1. A method for automatic segmentation of brain tumor images, characterized in that: include: Build the first brain tumor segmentation model based on Pytorch coding and nnUNet framework; Converting the first brain tumor segmentation model from a Pytorch format to an ONNX format to obtain a second brain tumor segmentation model; Visualizing the second brain tumor segmentation model, and based on the visualized model structure, pruning the second brain tumor segmentation model to obtain a third brain tumor segmentation model; The third brain tumor segmentation model is quantized by the TensorRT optimization tool to obtain a brain tumor image automatic segmentation model for automatic brain tumor image segmentation; in: The first brain tumor segmentation model is constructed based on Pytorch coding and nnUNet framework, including: Using Pytorch coding, image preprocessing and image enhancement operations were performed on brain MRI images and corresponding images labeled with brain tumor labels to obtain 3D image data; Performing three-dimensional sliding window clipping on the 3D image data to obtain training samples; Providing a nnUNet network model, and training the nnUNet network model using the training samples; Performing a preliminary pre-training loss analysis in the nnUNet network model according to the training samples, selecting a corresponding network depth and number of convolution kernels for losses in different ranges and downward trends, and obtaining an adaptive network parameter configuration; By configuring the adaptive network parameters, relatively feasible network hyperparameters are obtained, and the nnUNet network model is optimized according to the network hyperparameters to obtain a brain tumor segmentation network model; Using brain MRI images as input of the brain tumor segmentation network model, inferring through the brain tumor segmentation network model, performing noise processing on the inference results to obtain predicted brain tumor labels, testing the brain tumor segmentation network model, and obtaining a first brain tumor segmentation model in Pytorch format; The step of visualizing the second brain tumor segmentation model and pruning the second brain tumor segmentation model based on the visualized model structure to obtain a third brain tumor segmentation model includes: Using the Netron visualization tool, generating a visualization model structure for the second brain tumor segmentation model; Confirming useless outputs in the visualization model structure, and obtaining node names, weight names and / or serial numbers corresponding to the useless outputs; Traversing and deleting the useless outputs in the second brain tumor segmentation model by using the node names, weight names and / or serial numbers; The deleted model is saved to obtain the third brain tumor segmentation model.

2. The method for automatic segmentation of brain tumor images according to claim 1, characterized in that: The first brain tumor segmentation model is converted from the Pytorch format to the ONNX format to obtain a second brain tumor segmentation model, including: Acquire the network structure and network parameters of the first brain tumor segmentation model; The network structure and the network parameters are passed into the torch.onnx.export model to obtain a second brain tumor segmentation model in ONNX format; The dynamic_axes parameter is added to the second brain tumor segmentation model in the ONNX format to set the input and output to dynamic dimensions so that the shape and size of the input and output can change in different inference requests.

3. The method for automatic segmentation of brain tumor images according to claim 1, characterized in that: The method of quantizing the third brain tumor segmentation model through the TensorRT optimization tool to obtain an automatic brain tumor image segmentation model includes: The third brain tumor segmentation model in ONNX format is converted into a TensorRT format model by using the trtexec program of the TensorRT optimizer, and quantization parameters are selected for the TensorRT format model to obtain a brain tumor image automatic segmentation model in TensorRT format.

4. The method for automatic segmentation of brain tumor images according to claim 3, characterized in that: Also includes any one or more of the following: - the quantization parameter selection includes: selecting fp16 or int8 as the quantization parameter; - The method uses the trtexec program of the TensorRT optimizer to convert the third brain tumor segmentation model in ONNX format into a model in TensorRT format, and selects quantization parameters for the model in TensorRT format to obtain an automatic brain tumor image segmentation model in TensorRT format, and also includes: Use the trtexec program of the TensorRT optimizer to build a TensorRT reasoning class, which is used to encapsulate the TensorRT model segmentation reasoning process; Deserialize a TensorRT engine from the given ONNX format model file path of the third brain tumor segmentation model, create the TensorRT engine execution context, and allocate host memory and device memory for input and output of the TensorRT engine; Define a memory allocation method and an inference execution method; wherein: the memory allocation method is used to allocate input and output host memory and device memory and related buffers for the TensorRT engine, and traverse all bindings of the TensorRT engine to allocate corresponding host memory and device memory according to the input or output type of the binding, and add the device memory to the binding list; the inference execution method is used to execute the segmentation inference process, wherein the segmentation inference process includes: first setting the input data to the passed array and executing the TensorRT inference, then extracting the output image from the inference result and returning the inference result; Instantiate the TensorRT inference class, read in the quantized TensorRT model, and build a brain tumor image automatic segmentation model in TensorRT format; By using the brain tumor image automatic segmentation model in the TensorRT format, the brain MRI image to be segmented is input into the TensorRT reasoning class of the brain tumor image automatic segmentation model, and the segmentation reasoning process is executed to perform automatic segmentation of the brain tumor image.

5. The method for automatic segmentation of brain tumor images according to any one of claims 1 to 4, characterized in that: Also includes any one or more of the following: -Build a Docker container environment, which is used to provide a model building environment; -The Flask framework is used to encapsulate the brain tumor image automatic segmentation model.

6. The method for automatic segmentation of brain tumor images according to claim 5, characterized in that: The method of encapsulating the brain tumor image automatic segmentation model using the Flask framework includes: The brain tumor image automatic segmentation model is encapsulated by using a Flask framework to obtain a Flask application; Use Dockerfile to build an image, and start the container with the built image to expose the port number and increase the memory limit at startup; Open the GPU and automatically start the Flask application when the image startup container is started, which is used to externally call the corresponding Flask brain tumor image automatic segmentation service.

7. A brain tumor image automatic segmentation system, characterized in that: include: The Pytorch model module builds the first brain tumor segmentation model based on Pytorch coding and the nnUNet framework; An ONNX model module, the module being used to convert the first brain tumor segmentation model from a Pytorch format to an ONNX format to obtain a second brain tumor segmentation model; a pruning module, the module being used to visualize the second brain tumor segmentation model, and based on the visualized model structure, prune the second brain tumor segmentation model to obtain a third brain tumor segmentation model; A TensorRT model module, which is used to quantize the third brain tumor segmentation model through the TensorRT optimization tool to obtain a brain tumor image automatic segmentation model for automatic brain tumor image segmentation; in: The Pytorch model module, based on Pytorch coding and nnUNet framework, builds the first brain tumor segmentation model, including: Using Pytorch coding, image preprocessing and image enhancement operations were performed on brain MRI images and corresponding images labeled with brain tumor labels to obtain 3D image data; Performing three-dimensional sliding window clipping on the 3D image data to obtain training samples; Providing a nnUNet network model, and training the nnUNet network model using the training samples; Performing a preliminary pre-training loss analysis in the nnUNet network model according to the training samples, selecting a corresponding network depth and number of convolution kernels for losses in different ranges and downward trends, and obtaining an adaptive network parameter configuration; By configuring the adaptive network parameters, relatively feasible network hyperparameters are obtained, and the nnUNet network model is optimized according to the network hyperparameters to obtain a brain tumor segmentation network model; Using brain MRI images as input of the brain tumor segmentation network model, inferring through the brain tumor segmentation network model, performing noise processing on the inference results to obtain predicted brain tumor labels, testing the brain tumor segmentation network model, and obtaining a first brain tumor segmentation model in Pytorch format; The pruning module visualizes the second brain tumor segmentation model and prunes the second brain tumor segmentation model based on the visualized model structure to obtain a third brain tumor segmentation model, including: Using the Netron visualization tool, generating a visualization model structure for the second brain tumor segmentation model; Confirming useless outputs in the visualization model structure, and obtaining node names, weight names and / or serial numbers corresponding to the useless outputs; Traversing and deleting the useless outputs in the second brain tumor segmentation model by using the node names, weight names and / or serial numbers; The deleted model is saved to obtain the third brain tumor segmentation model.

8. The brain tumor image automatic segmentation system according to claim 7, characterized in that: Also includes any one or more of the following: - an environment building module, which is used to build a Docker container environment, and the Docker container environment is used to build the Pytorch model module, the ONNX model module, the pruning module and the TensorRT model module; - A Flask application module, which uses the Flask framework to encapsulate the brain tumor image automatic segmentation model.

9. A computer terminal comprising a memory, a processor and a computer program stored in the memory and executable on the processor, characterized in that: When the processor executes the computer program, it can be used to perform the method described in any one of claims 1 to 6, or to run the system described in any one of claims 7 to 8.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, it can be used to perform the method described in any one of claims 1 to 6, or to run the system described in any one of claims 7 to 8.

Images