Lightweight 3D medical image real-time reasoning method and system based on edge calculation

Through lightweight student model and CUDA parallelization technology, the problems of large size and low processing efficiency of nnUNet framework models are solved, efficient edge device deployment and real-time inference are achieved, and processing efficiency and accuracy are improved.

CN120339267AActive Publication Date: 2025-07-18XIAOZHI FUTURE (CHENGDU) TECH CO LTD

Patent Information

Application Number
CN202510788224.1
Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
Filing Date
2025-06-13
Publication Date
2025-07-18
Estimated Expiration
2045-06-13

AI Technical Summary

Technical Problem

The existing nnUNet framework model is large in size, has high requirements for computing GPUs, has low upper limit for large-scale data processing, low processing efficiency, and requires multiple model inferences in real-time auxiliary tasks, resulting in time loss.

Method used

A lightweight 3D medical image real-time inference method based on edge computing is adopted, a lightweight student model is built through a C++ framework, combined with knowledge distillation technology, the model volume is reduced to 1MB-3MB, and a CUDA parallelization technology is used to accelerate GPU, reduce the number of inferences, and optimize the sliding window inference process.

Benefits of technology

Significantly shortens the single inference time, reduces the model's memory occupancy, adapts to edge device deployment, maintains high accuracy, improves the performance of the entire process by more than 10 times, adapts to different hardware devices, supports 4D data processing and dynamic organ motion analysis.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN120339267A_ABST
    Figure CN120339267A_ABST
Patent Text Reader

Abstract

The invention relates to the technical field of image data processing, and particularly discloses a lightweight 3D medical image real-time reasoning method and system based on edge computing, and the method comprises the steps: building a C + + frame model, carrying out anisotropy detection and processing, separating low-resolution axis resampling, carrying out 3D resampling, determining the step length of a sliding window, and generating a weight map. And performing prediction, sliding window reasoning and image post-processing by using a lightweight student model, performing processing including size adjustment, voxel communication and integrity filling on a predicted segmentation map obtained through reasoning, and outputting a final global segmentation map. According to the method, the lightweight student model is adopted to replace the traditional nnUNet model reasoning which needs five times, the reasoning frequency is reduced by 80%, and the single reasoning time is remarkably shortened; meanwhile, CUDA parallelization is achieved through C + + language reconstruction, and through data block extraction, mirror image transformation and GPU acceleration of a weighted aggregation operator, the whole process performance is improved by more than 10 times compared with Python implementation.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the technical field of image data processing, and particularly to the real-time inference processing technologies such as the segmentation and marking of 3D medical images. Specifically, it relates to a lightweight 3D medical image real-time inference method and system based on edge computing. Background Art

[0002] The core of current medical image segmentation technology has shifted from traditional methods to deep learning-driven. Traditional techniques such as threshold segmentation, region growing, and atlas-based registration methods are still practical in specific scenarios (such as high-contrast organ segmentation), but generally face the problems of complex artificial feature design and weak generalization ability. In recent years, encoder-decoder architectures represented by U-Net have become mainstream, achieving multi-scale feature fusion through skip connections and making breakthroughs in tumor and organ segmentation tasks. Its improved versions such as Attention U-Net with attention mechanism added, three-dimensional extended V-Net, and models such as TransUNet combined with Transformer have significantly improved the localization accuracy of complex lesion areas.

[0003] The technological development shows three major trends: First, multi-modal fusion technology improves the segmentation robustness by integrating multi-source data such as CT, MRI, and PET and utilizing cross-modal feature complementarity; second, self-supervised / weakly-supervised learning strategies, such as contrast learning and label propagation, effectively alleviate the pain point of scarce labeled data in medical images; in addition, real-time interactive segmentation systems are gradually integrated into the clinical workflow. The current challenges focus on small-sample generalization ability, model interpretability verification, and domain adaptation problems under different devices / protocols. Federated learning, domain adaptation, and explainable AI have become the key breakthrough directions.

[0004] At present, the most widely used for medical image inference is the general medical processing framework nnUNet. As a commonly used framework for medical processing, nnUNet has been unanimously recognized in the industry for the flexibility and simplicity of its API. However, in terms of engineering and efficiency, the native nnUNet is not very developer-friendly during use due to the limitations of its overall architecture and the Python language of its API, and even sacrifices time cost due to accuracy. Its main limitations are reflected in that 3D model training requires high GPU memory, which has obvious limitations for large-scale data applications. At the same time, there may be additional fine-tuning for extreme data distributions. Summary of the Invention

[0005] To solve the problems existing in the existing nnUNet framework model, such as large model volume, high requirements for the computing GPU, low upper limit for processing large-scale data, and low processing efficiency, etc., this application specifically developed a lightweight 3D medical image real-time inference method and system based on edge computing to replace the existing nnUNet framework model. On the premise of basically retaining the accuracy of the existing nnUNet framework model, the size of the model was reduced from the existing volume of 100MB - 200MB to 1MB - 3MB, and the volume was basically reduced to about 1% - 2% of the existing nnUNet framework model, reducing the model occupancy by more than 98%. The existing nnUNet framework model is restricted by the original 5-fold model for separate inference and then aggregation prediction at the post-processing level. As a result, in order to obtain a higher prediction accuracy, 5 times of hardware time loss is required, which poses great difficulties for medical-related real-time assistance tasks.

[0006] Based on the general framework of nnUNet, this application combined the rich training framework and strategies of nnUNet itself to implement the knowledge distillation technology of the brand-new nnUNet based on the feature map layer by layer and progressively. This enables the present invention to seamlessly connect with the nnUNet framework, making the operation very easy. After testing, the Dice coefficient loss of the prediction results of the distilled student model and the 5 models of the original nnUNet is generally within 5%. If some hyperparameters of the distillation architecture are changed, the accuracy of this model is even better than the effect of the original 5 models for separate prediction and then aggregation. At the same time, this is also the fundamental reason why the model volume of the present invention is much smaller than the existing nnUNet framework model.

[0007] To achieve the above objectives, the technical solutions adopted in this application are as follows: The lightweight 3D medical image real-time inference method based on edge computing provided by the present invention includes the following steps: Step STP100, building the C++ framework model, establishing a 5-fold teacher model and a lightweight student model based on the nnUNet model, using the 5-fold teacher model to predict the sample image data, and globally guiding the lightweight student model through the simultaneously output Logits-P and feature maps until the prediction accuracy of the lightweight student model reaches the preset value of the teacher model's prediction accuracy, and then completing the establishment of the lightweight student model; Step STP200, anisotropy detection and processing, performing anisotropy detection on the input 3D medical image, and judging whether the extreme value ratio of the resolution of any axis γ difference meets the threshold γ0. When the ratio γ > γ0, then proceed to step STP300; when the ratio γ ≤γ0, then proceed to the resampling strategy of step STP400; where the ratio γis the ratio of the maximum value to the minimum value of the image voxels on the corresponding axis; Step STP300, perform separate low-resolution axis resampling, check whether the shape of the original 3D medical image is the same as that of the target 3D medical image. If they are the same, execute Step STP500. If not, use OpenCV to perform independent interpolation 2D resampling on the 2D plane perpendicular to the low-resolution axis, and then perform 1D resampling on the result of the 2D resampling in the low-resolution axis direction to complete 2D+1D resampling; Step STP400, 3D resampling. First, use the OpenCV function to perform 2D resampling on each XY plane perpendicular to the low-resolution axis, then collect the low-resolution axis pixels at the same (y, x) positions to form 1D lines, and then perform resampling on these 1D lines in the low-resolution axis direction using the resize of OpenCV; Step STP500, determine the sliding window step size and generate a weight map. Use the window step size calculator to automatically determine the overlapping step size according to the input volume size and block size, and use the Gaussian weighted aggregator to generate a weight map by weight assignment using the pre-computed 3D Gaussian kernel ; Step STP600, use the lightweight student model for prediction, and use the student model to predict the probability distribution Logits-C of each voxel point of the 3D medical image; Step STP700, sliding window inference. Use the CUDA operator module to perform block extraction, overlapping sliding, enhancement, aggregation, and normalization processing on the input 3D medical image to obtain the predicted global prediction segmentation map; Step STP800, image post-processing. Perform processing including size adjustment, voxel connectivity, and integrity filling on the predicted segmentation map obtained by inference in Step STP700, and output the final global segmentation map.

[0008] Preferably, the steps of block extraction and overlapping sliding of the CUDA operator module on the input 3D medical image in Step STP700 include Step STP710: the volume of the large-size 3D medical image is decomposed into a series of blocks of a fixed size window through the CUDA kernel function , and the possibly overlapping image blocks are denoted as . The window is slid in three dimensions of Z, Y, and X on the GPU in the form of a matrix at a predefined step size, and the starting coordinates of each block are generated ([[]] d, h, w ), which is expressed as: where, , , , and , The image patch representing the sliding window volume size is the streaming form of the CUDA kernel.

[0009] Further preferably, the step of the CUDA operator module in step STP700 for enhancing the input 3D medical image includes step STP720: For each extracted image patch Apply a set of predefined spatial transformations processed by CUDA kernel functions on the GPU, perform inference on each transformed patch according to the CUDA stream, then perform the inverse transformation on the CUDA kernel for the obtained prediction results, and finally average the prediction results of all transformed versions on the GPU, specifically expressed as: Among them, is the image patch The average predicted Logits after TTA, respectively represent the depth, height, and width dimension indices, with the size of , is the number of classes, is the streaming form of the CUDA kernel; represents the set of predefined spatial transformations, represents the set A specific transformation index in; represents the th spatial transformation function, is to The operation of restoring the predicted result after transformation to the original patch direction, that is, the inverse function of the th transformation; represents the inference function, and the spatial transformation includes combinations of mirror flips and / or non-flips along the Z, Y, and X axes.

[0010] Further preferably, the step of the CUDA operator module in step STP700 for aggregating the input 3D medical image includes step STP730: Specifically, the Gaussian weighted aggregation method is adopted, which specifically includes introducing a pre-computed Gaussian weight map with the same size as the patch in the GPU cache area , and then for each image patch The predicted result averaged through step STP720 When adding to the global prediction atlas, perform element-wise weighting on it using the CUDA Gaussian kernel function on the GPU to obtain the global cumulative probability distribution, specifically expressed as: Among them, is the Gaussian weight map, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel, is the GPU cache of the CUDA kernel function.

[0011] Still further preferably, the step of the CUDA operator module normalizing the input 3D medical image in step STP700 includes step STP740: After the weighted predictions of all image patches are accumulated into the global GPU cache, the final prediction of each voxel is obtained by dividing the weighted Logits accumulated for each voxel by the total weight accumulated for each voxel , specifically expressed as: wherein, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel.

[0012] The present invention also provides a lightweight 3D medical image real-time inference system based on edge computing for performing the real-time inference method provided above, including a feature distillation module and an inference engine module based on C++. The feature distillation module includes distilling the 5-fold teacher model set of the original nnUNet into a single lightweight student model, and the volume of the student model is 1MB - 3MB with an accuracy loss of no more than 5%; the inference engine module is used to achieve full-process acceleration from model conversion, preprocessing, sliding window inference to post-processing, and supports CUDA parallel computing; a sliding window optimization module, which adopts a GPU-driven overlapping window inference technology to achieve smooth output of prediction results through Gaussian weighted aggregation; an edge deployment interface module, which adapts to the personalized deployment requirements of different hardware devices.

[0013] Preferably, the feature distillation module further includes a progressive knowledge transfer unit for optimizing the parameters of the student model by comparing layer-by-layer feature maps and global guidance during the training stage of the student model. The progressive knowledge transfer unit includes an adaptive parameter module for dynamically adjusting the knowledge transfer path and weight, a feature map distillation module for aligning the intermediate semantic expression feature maps of the teacher model and the student model, and a self-supervised contrast module for enhancing the representation by learning the relationship between samples.

[0014] Preferably, the inference engine module includes an anisotropic processor for detecting the physical spacing ratio of the medical image and starting the separated Z-axis processing when the following formula is satisfied: wherein, Represents the extreme ratio of any axis resolution, Resolution extreme ratio threshold, Is the maximum resolution value, Is the minimum resolution value; The sliding window optimization module includes a window step calculator that automatically determines the overlapping step according to the input volume size and block size: Among them, Is the image block size, = 0.5, indicating 50% overlap; CUDA stream controller, which manages the asynchronous pipelined execution of data block extraction, mirror enhancement, inference, and aggregation; Gaussian weighted aggregator, which uses a pre-computed 3D Gaussian kernel for weight assignment: Among them, Is the Gaussian weight map, Is the in-block coordinate, Is the block center coordinate, Is the block in dimension i The size of, standard deviation Is set to be proportional to the block size, = 0.5.

[0015] The present invention also provides a computer-readable storage medium, on which a computer program is stored. When the computer program is executed by a processor, the steps of the real-time inference method described above are implemented.

[0016] The present invention also provides a medical image processing device, including a memory, a processor, and a computer program stored on the memory. When the processor executes the program, the steps of the above real-time inference method are implemented.

[0017] Beneficial effects: 1. The present invention uses a single lightweight student model to replace the traditional nnUNet that requires 5 model inferences, reducing the number of inferences by 80% and significantly shortening the single-inference time; at the same time, it is reconstructed using C++ language to achieve CUDA parallelization. Through the GPU acceleration of data block extraction, mirror transformation, and weighted aggregation operators, the full-process performance is improved by more than 10 times compared with the Python implementation.

[0018] 2. The present invention adopts model lightweighting and adapts to edge device deployment, overcoming the difficulty of deploying existing large-size models of hundreds of MB in edge devices with limited video memory. Through feature map migration and channel pruning, the model volume is compressed to 1-3MB. At the same time, Python dependencies are removed, and the video memory occupancy is reduced by 70%, enabling deployment to embedded GPUs or medical terminal devices.

[0019] 3. The present invention dynamically calibrates the intermediate layer features of the student model by the teacher model, retains dark knowledge, controls the Dice coefficient loss to ≤5%, and the accuracy loss is controllable. Moreover, the prediction results in the overlapping area of the sliding window are smoothed and weighted to avoid the accumulation of edge prediction errors, improve the continuity of the segmentation boundary, and make the segmentation results reliable.

[0020] 4. Compared with the problem that direct 3D resampling in the existing nnUNet framework model is prone to loss of low-resolution axis information or artifacts, the present invention dynamically selects a 2D+1D resampling strategy according to the physical spacing ratio to retain the details of each axis; after interpolation, it is cropped to the original data range to avoid the interference of outliers.

[0021] 5. The present invention is reconstructed using a brand-new C++ language, can be compiled for ARM / x86 architectures, and adapts to full-scenario deployment from cloud servers to portable ultrasound devices. Automatically adjusts the sliding window batch size according to the device computing power, and still maintains real-time performance on low-end devices; modular CUDA operator library, supporting user-defined processes.

[0022] 6. The present invention can be extended to 4D data processing (3D space + time dimension), supports dynamic organ motion analysis, and the preprocessing module automatically adapts to the physical characteristics of different modality images such as CT, MRI, and PET, specifically solving the problem of insufficient processing ability of existing general frameworks for 4D images, such as cardiac dynamic MRI. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the prior art, the following will briefly introduce the drawings required for use in the description of the embodiments or the prior art. Obviously, the following drawings are only some embodiments of the present application. For those skilled in the art, without creative efforts, other drawings can be obtained based on these drawings.

[0024] Figure 1 It is a schematic block diagram of the system data flow during the system training and use stages of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0025] To make the objectives, technical solutions, and advantages of the embodiments of this application clearer, the following will clearly and completely describe the technical solutions in the embodiments of this application with reference to the accompanying drawings in the embodiments of this application. Obviously, the described embodiments are some, but not all, of the embodiments of this application. Components of the embodiments of this application described and illustrated in the accompanying drawings here can be arranged and designed in various different configurations.

[0026] Therefore, the following detailed description of the embodiments of this application provided in the accompanying drawings is not intended to limit the scope of this application claimed, but merely represents selected embodiments of this application. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of this application without creative efforts fall within the scope of protection of this application.

[0027] Embodiment 1: The lightweight 3D medical image real-time inference method based on edge computing provided by the present invention includes the following steps: Step STP100, building a C++ framework model, establishing a 5-fold teacher model and a lightweight student model based on the nnUNet model, using the 5-fold teacher model to predict sample image data, and globally guiding the lightweight student model through the simultaneously output Logits-P and feature maps until the prediction accuracy of the lightweight student model reaches a preset value of the prediction accuracy of the teacher model, and then completing the establishment of the lightweight student model; Step STP200, anisotropy detection and processing, performing anisotropy detection on the input 3D medical image, and judging whether the ratio of the extreme value of the resolution of any axis γ difference meets the threshold γ0. When the ratio γ > γ0, then go to step STP300; when the ratio γ ≤ γ0, then go to step STP400 resampling strategy; where the ratio γ is the ratio of the maximum value to the minimum value of the image voxels on the corresponding axis; Step STP300, separating the low-resolution axis resampling, checking whether the shapes of the original 3D medical image and the target 3D medical image are the same. If they are the same, then execute step STP500. If they are different, then use OpenCV to perform independent interpolation 2D resampling on the 2D plane perpendicular to the low-resolution axis, and then perform 1D resampling on the result of the 2D resampling in the low-resolution axis direction to complete 2D+1D resampling; Step STP400, 3D resampling. First, use OpenCV functions to perform 2D resampling on each XY plane perpendicular to the low-resolution axis, then collect the low-resolution axis pixels at the same (y, x) positions to form 1D lines, and then use OpenCV's resize to perform resampling in the low-resolution axis direction on these 1D lines; Step STP500, determine the sliding window step size and generate a weight map. Use a window step size calculator to automatically determine the overlapping step size according to the input volume size and block size, and use a Gaussian weighted aggregator to generate a weight map by weight distribution using a pre-computed 3D Gaussian kernel ; Step STP600, use a lightweight student model for prediction. Use the student model to predict the probability distribution Logits-C of each voxel point in the 3D medical image; Step STP700, sliding window inference. Use the CUDA operator module to perform block extraction, overlapping sliding, enhancement, aggregation, and normalization on the input 3D medical image to obtain a predicted global prediction segmentation map; Step STP800, image post-processing. Perform processing including size adjustment, voxel connectivity, and integrity filling on the predicted segmentation map obtained through inference in Step STP700, and output the final global segmentation map.

[0028] Example 2: This example further optimizes Step STP700 on the basis of Example 1. Specifically, the steps of the CUDA operator module in this example for performing block extraction and overlapping sliding on the input 3D medical image include Step STP710: The volume of the large-size 3D medical image is decomposed into a series of blocks of a fixed size window through a CUDA kernel function , and the possibly overlapping image blocks are denoted as . The window slides in the three dimensions of Z, Y, and X on the GPU in the form of a matrix at a predefined step size, generating the starting coordinates of each block ( d, h, w ), which is expressed as: where, , , , and , represents the image block representing the sliding window volume size, is the stream form of the CUDA kernel.

[0029] In this embodiment, the steps of the CUDA operator module for enhancing the input 3D medical image in step STP700 include step STP720: for each extracted image patch Apply a set of predefined spatial transformations processed by CUDA kernel functions on the GPU, perform inference on each transformed patch according to the CUDA stream, then perform an inverse transformation on the obtained prediction results on the corresponding CUDA kernel, and finally average the prediction results of all transformed versions on the GPU, which is specifically expressed as: Among them, is the average predicted Logits of the image patch after TTA, represent the depth, height, and width dimension indices respectively, with the size of , is the number of classes, is the stream form of the CUDA kernel; represents the set of predefined spatial transformations, represents the set a specific transformation index in; represents the th spatial transformation function, is the operation of restoring the prediction result after transforming to the original patch direction, that is, the inverse function of the th transformation; represents the inference function, and the spatial transformation includes combinations of mirror flips and / or non-flips along the Z, Y, and X axes.

[0030] In this embodiment, the steps of the CUDA operator module for aggregating the input 3D medical image in step STP700 include step STP730: specifically adopt the Gaussian weighted aggregation method, which specifically includes introducing a pre-computed Gaussian weight map with the same size as the block size in the GPU cache area , and then when adding the prediction result averaged by step STP720 of each image patch to the global prediction map, perform element-wise weighting on it using the CUDA Gaussian kernel function on the GPU to obtain the global cumulative probability distribution, which is specifically expressed as: Among them, is the Gaussian weight map, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel, It is the GPU cache of the CUDA kernel function.

[0031] In this embodiment, the steps of the CUDA operator module normalizing the input 3D medical image in step STP700 include step STP740: After the weighted predictions of all image patches are accumulated into the global GPU cache, the final prediction for each voxel is obtained by dividing the weighted Logits accumulated for each voxel by the total weight accumulated for each voxel , specifically expressed as: where, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel.

[0032] Embodiment 3: This embodiment of the present invention provides a lightweight 3D medical image real-time inference system based on edge computing for performing the real-time inference method provided in any of the above embodiments, including a feature distillation module and an inference engine module based on C++. The feature distillation module includes distilling the 5-fold teacher model set of the original nnUNet into a single lightweight student model, and the volume of the student model is 1MB - 3MB and the accuracy loss does not exceed 5%; the inference engine module is used to achieve full-process acceleration from model conversion, preprocessing, sliding window inference to post-processing, supporting CUDA parallel computing; a sliding window optimization module, which adopts a GPU-driven overlapping window inference technology to achieve smooth output of prediction results through Gaussian weighted aggregation; an edge deployment interface module, which adapts to the personalized deployment requirements of different hardware devices. The system framework principle is shown in Figure 1 as shown. Among them, for the native nnUNet, it includes a preprocessing module - P, an empirical parameter module, a cross-validation module, a fixed parameter module, and a post-processing module - P, and then outputs from the native nnUNet framework to the 5-fold teacher model.

[0033] In this embodiment, the feature distillation module further includes a progressive knowledge transfer unit for optimizing the parameters of the student model by comparing layer-by-layer feature maps and global guidance during the training stage of the student model. The progressive knowledge transfer unit includes an adaptive parameter module for dynamically adjusting the knowledge transfer path and weight, a feature map distillation module for aligning the intermediate semantic expression feature maps of the teacher model and the student model, and a self-supervised contrast module for enhancing the representation by learning the relationship between samples.

[0034] In this embodiment, the inference engine module includes an anisotropic processor for detecting the physical spacing ratio of the medical image and starting the separated Z-axis processing when the following formula is satisfied: Among them, represents any axis resolution extreme ratio, resolution extreme ratio threshold, is the maximum resolution, is the minimum resolution; The sliding window optimization module includes a window step calculator that automatically determines the overlap step according to the input volume size and block size: Among them, is the image block size, = 0.5, indicating 50% overlap; CUDA stream controller, which manages the asynchronous pipelined execution of data block extraction, mirror enhancement, inference, and aggregation; Gaussian weighted aggregator, which uses a pre-computed 3D Gaussian kernel for weight assignment: Among them, is the Gaussian weight map, are the coordinates within the block, is the center coordinate of the block, is the size of the block in dimension i , and the standard deviation is set to be proportional to the block size, = 0.5.

[0035] Regarding the implementation of CUDA operators in the C++ framework, the following explanations are made in this embodiment: First, utilize the parallel computing power of the graphics processing unit (GPU) to efficiently implement CUDA operators for the key computational steps in the sliding window inference process. These operators are designed to accelerate dense prediction tasks for large-scale volume data, such as 3D medical images, by offloading computationally intensive operations from the CPU to the GPU and leveraging the CUDA architecture to achieve large-scale parallelization, thereby significantly enhancing the processing speed and efficiency. These operators cover the complete process from data block extraction, spatial transformation, weighted aggregation to the integration of the final result, and provide multiple optimized implementation methods such as single-block, streaming asynchronous, and batch processing.

[0036] Data block extraction operator: Efficiently copy a rectangular sub-region (Patch) with specified starting coordinates and dimensions from the large 3D input image tensor stored in the GPU video memory as the input for subsequent model inference.

[0037] Using CUDA threads, each thread is responsible for calculating and copying one or more voxel values in the target Patch buffer. For each thread, a one-dimensional index in the target Patch buffer is calculated based on its unique thread ID. The coordinates are decomposed into three-dimensional coordinates within the Patch. Combined with the starting coordinates (d, h, w) of the Patch in the original image, the three-dimensional coordinates of the voxel in the source image buffer are calculated, and the three-dimensional source image coordinates are converted into a one-dimensional index in the source image buffer, and read from the value, and write it into .

[0038] The kernel is designed to handle multiple Patches in a single batch. The thread index is further decomposed to include the item index within the batch. The starting coordinates (d, h, w) are read from the coordinate array in the GPU video memory, and each batch item corresponds to different coordinates. The kernel will handle potential padding items, and when the batch size cannot be evenly divided by the actual workload, the effective range is controlled by the size of the current batch.

[0039] Mirror transformation operator: Perform a mirror operation on a data block in the GPU video memory along one or more specified coordinate axes, where the depth Z = 0, the height Y = 1, and the width X = 2, as part of data augmentation TTA during testing or as part of the inverse transformation of the result.

[0040] Adopt a large-scale parallel strategy, where each thread is responsible for processing one element in the source data block and writing it to the mirrored position in the target data block. The thread calculates the one-dimensional index of the source element it is responsible for in the data block and decomposes it into three-dimensional coordinates For the mirror output pointer, it also includes the category index . According to the axis to be mirrored, calculate the target three-dimensional coordinates . For example, if axis = 0, then , and the other coordinates remain unchanged. Read the value from and write it into .

[0041] For the Host-side logic, since the mirror operation order is relevant, the host-side wrapper functions mirror*Cuda and mirrorBatch*Cuda usually adopt the "ping-pong buffer" strategy, using a temporary GPU buffer to apply the mirror kernel functions for each axis in sequence to ensure the correctness of multi-axis mirroring. At the same time, the kernel processes batch data and performs mirror operations independently on each Patch in the batch.

[0042] Weighted accumulation operator: Multiply the Logits of the predicted Patches after a single or a batch of inferences by the corresponding pre-computed Gaussian weight map, and atomically accumulate the weighted results into the global (covering the entire input volume size) predicted Logits buffer and the weight accumulation buffer.

[0043] Start a thread that covers all elements of the predicted Patch, including the class dimension. Each thread calculates the class it is responsible for and the spatial position pos within the Patch (decomposed into ). Combine the starting coordinates d, h, w of the Patch to calculate the coordinates of this voxel in the global volume . Read the predicted value P = pred[thread globalidx and the corresponding Gaussian weight , and calculate the weighted value .

[0044] Calculate the one-dimensional indices global logits in the global Logits buffer predicted predictions and the global weight buffer n logitsidx and globa lspatialidx .

[0045] Add safely to predicted logits [global logitsidx , and add safely to n predictions [global spatialidx . atomicAdd ensures that even if multiple threads, i.e., threads processing different Patches in the overlapping region, try to update the same memory location simultaneously, the result is correct, serializing access. The weight is usually only accumulated once by the thread with c = 0 to avoid redundancy. The kernel processes batch data, and the thread calculates batch idx , and reads from the GPU coordinate array (d, h, w) . Optimize the memory access pattern.

[0046] Fusion operator: Reduce the CUDA kernel launch overhead and GPU memory bandwidth occupancy by combining / fusing multiple logically consecutive steps in the sliding window process into a single kernel function for execution.

[0047] The data extraction and mirroring operations are completed within a single kernel function. The thread reads the source image data, directly calculates the coordinates within the mirrored target Patch, and writes the data to the final position in the target Patch buffer. The reverse mirroring and weighted accumulation operations of the model output are completed within a single kernel function. The thread reads the un-reversed model output value, calculates the global target coordinates after reverse mirroring, reads the corresponding Gaussian weight, calculates the weighted value, and then atomically accumulates it to the target position in the global buffer. This reduces the number of reads and writes of intermediate data, reduces the kernel function startup latency, may improve data locality, and thus enhances the overall performance.

[0048] Final aggregation operator: After the weighted prediction accumulation of all Patches is completed, the accumulated Logits are divided by the accumulated weights for final normalization to obtain the average predicted Logits for each voxel, and the data layout in the video memory may be adjusted.

[0049] Start threads that cover the entire volume space dimension. Each thread is responsible for a spatial position. Read the accumulated weight W = n predictions [spatial idx . If is greater than a small threshold to avoid division by zero, then calculate . Then, for all classes , read the accumulated Logit L = predicted logits [c × total voxels + spatial idx , calculate the average value , and write the result back to the predicted logits buffer to achieve in-place modification.

[0050] If it is necessary to change the memory layout of the final output, for example, from C, D, H, W to D, H, W, C, then start a parallel copy kernel function, calculate the indices according to the source and target layouts, and move the data. If the layout remains unchanged, no operation is required.

[0051] Asynchronous stream operator: It realizes the overlapping execution between the computing tasks and data transfer or other computing tasks, improving the GPU utilization rate. The *Parallel version of the above operator or the version with the cudaStream t parameter allows the submission of kernel function launches and memory operations to a specific CUDA stream. Combining with the CUDA event cudaEvent t for synchronization control, a complex execution pipeline can be constructed. For example, perform the inference of the current batch on one stream, while perform the Patch extraction and mirroring of the next batch on another stream, thus hiding part of the latency.

[0052] Generally speaking, the comparison of the effects of this embodiment with the existing native nnUNet is shown in Table 1 below: Table 1 Comparison of the effects between this embodiment and the native nnUNet framework model Example 4: The present invention also provides a computer-readable storage medium, on which a computer program is stored. When the computer program is executed by a processor, the steps of the real-time inference method described in Embodiments 1-2 above are implemented.

[0053] Example 5: The present invention also provides a medical image processing device, including a memory, a processor, and a computer program stored on the memory. When the processor executes the program, a lightweight 3D medical image real-time inference method based on edge computing described in Embodiments 1-2 above is implemented.

[0054] The above are only the preferred embodiments of the present application and are not intended to limit the present application. For those skilled in the art, the present application may have various changes and modifications. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present application shall be included in the protection scope of the present application.

Claims

1. A lightweight 3D medical image real-time inference method based on edge computing, characterized in that, It includes the following steps: Step STP100: Build a C++ framework model. Based on the nnUNet model, establish a 5-fold teacher model and a lightweight student model. Use the 5-fold teacher model to predict the sample image data, and use the simultaneously output Logits-P and feature maps to globally guide the lightweight student model until the prediction accuracy of the lightweight student model reaches the preset value of the teacher model's prediction accuracy, then complete the establishment of the lightweight student model; Step STP200, anisotropy detection and processing, perform anisotropy detection on the input 3D medical image by judging the extreme value ratio of the resolution of any axis γ whether the difference meets the threshold γ >, then proceed to step STP300; when the ratio γ ≤ then proceed to step STP400 resampling strategy; where the ratio γ is the ratio of the maximum value to the minimum value of the image voxels on the corresponding axis; Step STP300: Separate low-resolution axis resampling. Check whether the shape of the original 3D medical image is consistent with the shape of the target 3D medical image. If they are consistent, execute Step STP500. If not, use OpenCV to perform independent interpolation 2D resampling on the 2D plane perpendicular to the low-resolution axis, and then perform 1D resampling on the result of the 2D resampling in the low-resolution axis direction to complete 2D+1D resampling; Step STP400: 3D resampling. First, use the OpenCV function to perform 2D resampling on each XY plane perpendicular to the low-resolution axis, then collect the low-resolution axis pixels at the same (y, x) positions to form 1D lines, and then use the resize of OpenCV to perform resampling in the low-resolution axis direction on these 1D lines; Step STP500, determine the sliding window step size and generate a weight map. Use a window step calculator to automatically determine the overlapping step size according to the input volume size and block size, and use a Gaussian weighted aggregator to generate a weight map by performing weight assignment using a pre-computed 3D Gaussian kernel ; Step STP600: Use the lightweight student model for prediction. Use the lightweight student model to predict the probability distribution Logits-C of each voxel point of the 3D medical image; Step STP700: Sliding window inference. Use the CUDA operator module to perform block extraction, overlapping sliding, enhancement, aggregation, and normalization processing on the input 3D medical image to obtain the predicted global prediction segmentation map; Step STP800: Image post-processing. Perform processing including size adjustment, voxel connectivity, and integrity filling on the predicted segmentation map obtained by inference in Step STP700, and output the final global segmentation map.

2. The lightweight 3D medical image real-time inference method based on edge computing according to claim 1, wherein, In step STP700, the steps of block extraction and overlapping sliding of the input 3D medical image by the CUDA operator module include step STP710: the volume of the large-size 3D medical image is decomposed into a series of blocks of a fixed size window through the CUDA kernel function , and the possibly overlapping image blocks are denoted as . The window is slid in three dimensions of Z, Y, and X on the GPU in the form of a matrix according to the CUDA stream with a predefined step size to generate the starting coordinates of each block ( d, h, w ), which is expressed as: Among them, , , , and , an image block representing the volume size of the sliding window, is the stream form of the CUDA core.

3. The lightweight 3D medical image real-time inference method based on edge computing according to claim 2, wherein The steps for the CUDA operator module to enhance the input 3D medical image in step STP700 include step STP720: for each extracted image patch Apply a set of predefined spatial transformations processed by CUDA kernel functions on the GPU, perform inference on each transformed patch according to the CUDA stream, then perform the inverse transformation on the CUDA kernel for the obtained prediction results, and finally average the prediction results of all transformed versions on the GPU, which is specifically expressed as: Among them, is the average predicted Logits of the image patch after TTA, represent the depth, height, and width dimension indices respectively, with the size of , is the number of classes, is the stream form of the CUDA kernel; represents a predefined set of spatial transformations, represents a specific transformation index in the set; represents the th spatial transformation function, is the operation of restoring the predicted result after transforming to the original block direction, that is, the inverse function of the th transformation; represents the inference function, and the spatial transformations include combinations of mirror flips and / or non - flips along the Z, Y, and X axes.

4. The lightweight 3D medical image real-time inference method based on edge computing according to claim 3, wherein The steps of the CUDA operator module aggregating the input 3D medical image in step STP700 include step STP730: Specifically, the Gaussian weighted aggregation method is adopted, which specifically includes introducing a pre-computed Gaussian weight map with the same size as the block size in the GPU cache area , and then each image block The prediction result averaged through step STP720 When adding it to the global prediction map, it is weighted element by element on the GPU using the CUDA Gaussian kernel function to obtain the global cumulative probability distribution, which is specifically expressed as: Among them, is the Gaussian weight map, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel, is the GPU cache of the CUDA kernel function.

5. The lightweight 3D medical image real-time inference method based on edge computing according to claim 3, wherein The steps of the CUDA operator module normalizing the input 3D medical image in step STP700 include step STP740: all image patches After the weighted predictions are accumulated into the global GPU cache, the final prediction for each voxel is obtained by dividing the weighted Logits accumulated for each voxel by the total weight accumulated for each voxel , specifically expressed as: Among them, is the global cumulative probability distribution, is the global cumulative weight, initialized to 0; is the stream form of the CUDA kernel.

6. A lightweight 3D medical image real-time inference system based on edge computing, for performing the method according to any one of claims 1-5, characterized in that, It includes a feature distillation module and an inference engine module based on C++. The feature distillation module includes distilling the 5-fold teacher model set of the original nnUNet into a single lightweight student model. The volume of the lightweight student model is 1MB - 3MB and the accuracy loss does not exceed 5%. The inference engine module is used to achieve full-process acceleration from model conversion, preprocessing, sliding window inference to post-processing, and supports CUDA parallel computing. The sliding window optimization module adopts the GPU-driven overlapping window inference technology to achieve smooth output of the prediction results through Gaussian weighted aggregation. The edge deployment interface module adapts to the personalized deployment requirements of different hardware devices.

7. A lightweight 3D medical image real-time inference system based on edge computing according to claim 6, characterized in that, The feature distillation module further includes a progressive knowledge transfer unit for optimizing the parameters of the lightweight student model through layer-by-layer feature map comparison and global guidance during the lightweight student model training phase. The progressive knowledge transfer unit includes an adaptive parameter module for dynamically adjusting the knowledge transfer path and weights, a feature map distillation module for aligning the intermediate semantic expression feature maps of the teacher mode and the lightweight student model, and a self-supervised contrast module for enhancing the representation through inter-sample relationship learning.

8. A lightweight 3D medical image real-time inference system based on edge computing according to claim 6, characterized in that, The inference engine module includes an anisotropic processor for detecting the physical spacing ratio of medical images and starting the separated Z-axis processing when the following formula is satisfied: Among them, represents any axis resolution extreme value ratio, resolution extreme value ratio threshold, is the maximum resolution value, is the minimum resolution value; The sliding window optimization module includes a window stride calculator for automatically determining the overlapping stride according to the input volume size and block size: Among them, is the image block size, = 0.5, indicating a 50% overlap; A CUDA stream controller that manages the asynchronous pipelined execution of data block extraction, mirror enhancement, inference, and aggregation; A Gaussian weighted aggregator that performs weight assignment using a pre-computed 3D Gaussian kernel: Among them, is the Gaussian weight map, is the in-block coordinate, is the block center coordinate, is the size of the block in dimension i , and the standard deviation is set to be proportional to the block size, = 0.5.

Citation Information

Patent Citations

  • Interactive image editing

    CN114863076A

  • Medical image segmentation method fusing SAM global modeling and U-Net local optimization

    CN119832012A

  • Industrial product surface defect detection method based on AI large model

    CN120013927A

  • Image identification using neural networks

    US20200302176A1

  • System and method of training vision transformer on small-scale datasets

    US20240212330A1

Cited By

  • High-fidelity lightweight world model construction method for end-to-end automatic driving test

    CN120909949A

  • High-fidelity lightweight world model construction method for end-to-end autonomous driving test

    CN120909949B