SAM2 model-based zebra fish behavior video full-automatic segmentation method

Abstract

The invention discloses a zebra fish behavior video full-automatic segmentation method based on an SAM2 model, which comprises the SAM2 model, and is characterized in that a prompt automatic generation module and a memory bank detection optimization module are introduced on the basis of the SAM2 model, an image encoder adds Adapter for fine tuning, a mask decoder adds LoRA for fine tuning, and a memory bank detection optimization module is introduced on the basis of the SAM2 model. And the prompt automatic generation module inputs the generated prompt information into a mask decoder to obtain a full-automatic segmentation result of the zebra fish behavior video, and motion detection and segmentation efficiency optimization are performed through the memory bank optimization module. Compared with the prior art, the method has the advantages that only part of images need to be screened for labeling, video data does not need to be labeled frame by frame, the manual labeling cost is greatly reduced, the problem that the SAM2 model can generate a high-quality segmentation result only through high-quality prompting is well solved, and the method is simple, convenient, good in effect and high in efficiency. The method is especially suitable for real-time segmentation and tracking of a high-speed moving target in a microscope video, and has a good application prospect and commercial development value.

Description

Technical Field

[0001] The present invention relates to the technical field of biomedical image processing, and more particularly to a fully automatic segmentation method for zebrafish behavior videos based on the SAM2 model. Background Art

[0002] As an important model organism, zebrafish has unique advantages in the field of biomedical research. Its embryonic transparency allows direct observation of the organ development process, and its high genetic homology with humans (up to 87%) makes zebrafish an ideal model for drug toxicity testing, research on the mechanisms of neurodegenerative diseases, and screening of anti-cancer drugs. In phenotype-based drug screening, quantitative analysis of zebrafish behavior characteristics is a core indicator for evaluating drug effects, such as analysis of movement trajectories and modeling of social behaviors. The accurate analysis of these behavioral phenotypes highly depends on the pixel-level segmentation accuracy of the target object in video data. Traditional manual annotation methods have problems of low efficiency and subjective bias, which restrict the large-scale application of high-throughput drug screening. Biomedical image segmentation solutions based on deep learning can solve the problem of manual repetitive annotation, but the current mainstream solutions face several challenges: 1) The zebrafish tail has a high swinging frequency, and there will be a large amount of motion blur in the frame images obtained at a low shooting frame rate; 2) The semi-transparent characteristics of zebrafish lead to unclear edges and introduce the problem of background impurity interference; 3) The single-frame processing mode of the image segmentation solution loses the temporal information of motion, requires more redundant resources, and reduces the running efficiency.

[0003] DEVA decouples the video segmentation task into an image segmentation module and a bidirectional temporal propagation module, and fuses the multi-frame segmentation results through a task-agnostic propagation model to ensure cross-frame consistency. This module captures long-term temporal dependencies through bidirectional propagation (forward and backward), avoiding the cumulative errors of traditional unidirectional propagation. However, this solution has poor adaptability to the non-rigid motion of the zebrafish tail.

[0004] Segment Anything Model (SAM) has become a fundamental model in the field of image segmentation due to its powerful performance. Its subsequent work, the SAM2 model, performs excellently in the field of video segmentation by introducing a memory mechanism and an efficient Hiera image encoder. However, its drawback is that it requires manual given prompts, such as points, boxes, rough masks, etc., which is not conducive to deployment in fully automatic pipeline tasks.

[0005] PerSAM automatically generates point prompts based on the similarity relationship between images according to the maximum and minimum positions. Although this method is efficient, it is prone to target confusion when multiple corresponding targets appear in the same frame. Methods such as GroundedSAM introduce an additional object detection model and design a method for automatically generating box prompts for SAM. However, the introduction of the additional framework also greatly increases the resource consumption of model training and inference. SAM2Rad generates corresponding prompts using the image features obtained by the SAM image encoder, but does not use the key memory mechanism in SAM2 to generalize it to the video field.

[0006] In summary, the existing biomedical image processing cannot meet the real-time segmentation and tracking of high-speed moving targets in microscope videos, with relatively low segmentation accuracy and efficiency, and cannot obtain high-quality segmentation results. Summary of the Invention

[0007] The object of the present invention is to provide a fully automatic segmentation method for zebrafish behavior videos based on the SAM2 model in view of the deficiencies of the prior art. An automatic prompt generation module is introduced on the basis of the SAM2 model, and the prompt information generated by the automatic prompt generation module is input into the mask decoder. Combined with the memory bank optimization module, a high-quality fully automatic segmentation result of the zebrafish behavior video is efficiently obtained. This method generates the prompt information required by the SAM2 model for specific domain data and segmentation requirements. The automatic prompt generation module fuses multi-level image features and learnable prompt tokens through a bidirectional attention module and a feature pyramid network, and generates accurate prompt information using the prompt tokens during the inference process. The memory bank optimization module uses the correlation analysis of memory objects to detect the short-term high-speed tail-flicking motion of zebrafish, dynamically generates key frame prompts, and combines the memory information of past frames to improve the segmentation accuracy; for the long-term relatively static scenes of zebrafish, redundant memory objects are pruned to accelerate the calculation of the memory attention module and improve the segmentation efficiency. It preferably solves the problem that the SAM2 model needs high-quality prompts to produce high-quality segmentation results, and in scenarios where automation is required in practical applications, it is necessary to solve the problems of automatic prompt generation, when to give prompts, and the introduced additional computational consumption problems. The method is simple and effective, especially suitable for the real-time segmentation and tracking of high-speed moving targets in microscope videos, and has good application prospects and commercial development value.

[0008] The object of the present invention is achieved as follows: A fully automatic segmentation method for zebrafish behavior videos based on the SAM2 model, including all components of the SAM2 model. It is characterized in that a prompt automatic generation module and a memory bank detection optimization module are introduced on the basis of the SAM2 model. An Adapter is added to the image encoder for fine-tuning, and a LoRA is added to the mask decoder for fine-tuning. The prompt automatic generation module inputs the generated prompt information into the memory bank optimization module to obtain the fully automatic segmentation result of the zebrafish behavior video. The specific steps are as follows:

[0009] S1: Obtain the zebrafish behavior image dataset

[0010] Screen a partial set of frame images with strong motion features and impurity interference, and label the foreground of the zebrafish in the frame images to obtain a set of frame image labels. Among them, the artifacts generated by high-speed motion in the frame images are also considered as the foreground area and are also labeled with the foreground of the zebrafish;

[0011] S2: Construct and train the fully automatic segmentation model

[0012] Train the SAM2 model with the set of frame images and the set of frame image labels to obtain the trained model parameters, and construct the fully automatic segmentation model;

[0013] S3: Use the model to segment the zebrafish behavior video

[0014] Use the FFmpeg program to convert the zebrafish behavior video into a sequence of frame images, and input it into the fully automatic segmentation model loaded with the trained model parameters. Use the memory mechanism to propagate the historical frame information to obtain the fully automatic segmentation result of the zebrafish behavior video.

[0015] The prompt automatic generation module includes: learnable prompt tokens, a bidirectional attention module, a bounding box prediction multi-layer perceptron, and a rough mask prediction convolutional neural network; the memory bank detection optimization module includes: non-critical frame motion detection and redundant memory object pruning.

[0016] The learnable prompt tokens and the image feature embedding vectors generated by the image encoder and their corresponding position embedding vectors are input into the bidirectional attention module. The bidirectional attention module predicts the bounding box and the rough mask respectively through a multi-layer perceptron and a rough mask prediction convolutional neural network, and inputs them into the prompt encoder. The output of the prompt encoder is directly connected to the mask decoder as high-dimensional information.

[0017] The specific steps of training the SAM2 model in step S2 are as follows:

[0018] S2-1: Input the set of frame image samples into the fully automatic segmentation model, and obtain the prediction results of the bounding box and the rough mask through the image encoder and the prompt automatic generation module;

[0019] S2-2: Input the prediction results of the bounding box and the rough mask into the prompt encoder, and obtain the final predicted mask through the mask decoder;

[0020] S2-3: Calculate the generalized intersection over union loss through the predicted bounding box result and the ground truth bounding box generated from the frame image label;

[0021] S2-4: Calculate the focal loss through the predicted rough mask result and the ground truth mask of the frame image label;

[0022] S2-5: Calculate the Dice loss, focal loss, and boundary loss through the final predicted mask and the ground truth mask of the frame image label;

[0023] S2-6: Calculate the cross-entropy loss by predicting whether the decoder contains the target object and whether the ground truth mask contains the target object;

[0024] S2-7: Calculate the total loss by weighted calculation of all losses in steps S2-3 to S2-6;

[0025] S2-8: Use the AdamW optimizer to backpropagate the total loss to the fully automatic segmentation model and update all unfrozen parameters in the model;

[0026] S2-9: Repeat steps S2-1 to S2-8, and select the one with the highest Dice coefficient of the ground truth mask on the validation set as the optimal model parameters. The maximum number of repetitions of the above steps is 200 times.

[0027] The prompt automatic generation module includes: three bidirectional attention modules for fusing information of multi-level image features and learnable prompt tokens. Each bidirectional attention module has four layers. The first layer is a self-attention layer for the learnable prompt tokens to capture their own context information; the second layer is a token-image cross-attention layer for the learnable prompt tokens to extract information from the image features; the third layer is a multi-layer perceptron layer that performs non-linear transformation on the updated learnable prompt tokens through two fully connected layers; the fourth layer is an image-token cross-attention layer for injecting the information of the learnable prompt tokens back into the image features.

[0028] The image encoder Hiera is hierarchically designed. The features from low resolution to high resolution enter the bidirectional attention module for calculation in turn, serve as components of the query and key of the new layer, and fuse the multi-layer image features through the Feature Pyramid Network (FPN) to obtain the feature vector containing multi-scale information and the updated prompt tokens.

[0029] The feature vector containing multi-scale information passes through a lightweight mask hint prediction convolutional neural network to obtain a mask hint, and is used as a dense hint embedding vector through a hint encoder.

[0030] The first two labels of the updated hint tokens are used to predict the upper left and lower right coordinates of the bounding box hint, and a bounding box hint embedding vector is generated through the hint encoder. The remaining labels are concatenated after the bounding box hint embedding vector to form a sparse hint embedding vector.

[0031] The dense hint embedding vector and the sparse hint embedding vector are input into the mask decoder to obtain the final predicted mask, and the difference between it and the true label is calculated to obtain a loss function for optimizing the model parameters.

[0032] The generalized intersection over union loss function L of the bounding box box is calculated by the following formula:

[0033]

[0034] where A and B are the areas of the true bounding box and the predicted bounding box hint, respectively; C is the area of the smallest convex region containing A and B.

[0035] The focal loss Limask of the rough mask prediction result is calculated by the following formula:

[0036] L imask = L focal = -αy(1 - p) γ log(p) - (1 - α)(1 - y)p γ log(1 - p).

[0037] where y is the true mask, p is the output result of the rough mask prediction through the sigmoid function, and it is calculated pixel by pixel. α and γ are hyperparameters.

[0038] The final predicted mask loss L omask is calculated by the following formula:

[0039] L omask = λ focal L focal + λ dice L dice + λ bound L bound .

[0040] where L focal is calculated in the same way as the focal loss of the rough mask prediction result. y is the true mask, p is the final predicted mask of the model, and λ focal , λ dice and λ bound are all hyperparameters.

[0041] In addition, the calculation formula of L dice is expressed as follows:

[0042]

[0043] where P is the final predicted mask region and G is the ground truth mask region.

[0044] Finally, the calculation logic of L bound is to construct a normal vector for all points on the edge of the ground truth mask region and sum the L2 distances of the intersection segments with the edge of the predicted mask region. However, during the process of using it as a loss function, since the predicted mask region is indirectly represented by the softmax output and not explicitly defined, differential operations cannot be performed on the boundary points. Therefore, through mathematical derivation, it is transformed into the integral of the shortest signed distance from any point in the difference region to the ground truth mask edge, and the calculation formula is expressed as follows:

[0045] L bound = ∫ Ω φ G (q)s(q)dq.

[0046] where s(q) represents the predicted softmax value in the image region, and is the shortest signed distance from point q to the ground truth prediction edge.

[0047] The mask loss L that combines the final predicted mask loss and the rough mask prediction loss mask is calculated by the following formula:

[0048] L mask = λ i L imask + λ o L omask .

[0049] where λ i and λ o are both hyperparameters.

[0050] The loss L for whether the target object is contained obj is calculated by the following formula:

[0051] L obj = -ylog(p) - (1 - y)log(1 - p).

[0052] where y is the occurrence of the target object in the ground truth mask and p is the sigmoid value of the target object score output by the mask decoder.

[0053] The total loss function is calculated by the following formula:

[0054] Ltotal = λ1L box + λ2L mask + λ3L obj 。

[0055] Among them, λ1, λ2, and λ3 are all hyperparameters.

[0056] The specific steps of step S3 include:

[0057] S3-1: Convert the zebrafish behavior video data into a frame image sequence, input it into the trained fully automatic segmentation model, obtain all frame image feature embedding vectors through the image encoder, and mark the first frame as the key frame, and the rest are non-key frames;

[0058] S3-2: If the current frame is a non-key frame, the frame image feature embedding vector fuses the memory bank information through the memory attention,

[0059] and inputs it into the mask decoder to obtain the predicted mask; if the current frame is a key frame, the frame image feature embedding vector generates a prompt embedding vector through the prompt automatic generation module, and inputs it together with the image features containing memory information into the mask decoder to obtain the final predicted mask;

[0060] S3-3: For the predicted masks and original image features of all frames, generate memory objects through the memory encoder;

[0061] S3-4: If the current frame is a non-key frame, execute the non-key frame motion detection method. If it is re-marked as a key frame, return to step S3-2 to execute again, otherwise save the memory object to the nearest non-key nearest frame memory sequence and slide the memory window; if the current frame is a key frame, save the memory object to the key frame memory sequence;

[0062] S3-5: Execute the redundant memory object pruning method, and the pruned memory objects will not be used for the memory attention operation of the next frame;

[0063] S3-6: Perform the operations of S3-2 to S3-5 for each frame of the video. After all frames are processed, obtain the final zebrafish behavior video segmentation result.

[0064] The non-key frame motion detection and redundant memory object pruning methods adopt the following similarity calculation formula:

[0065]

[0066] where H and W are the height and width of the memory object respectively.

[0067] The calculation of the similarity is to flatten the spatial dimension of the memory object, calculate the cosine similarity position by position along the channel dimension in space, and finally calculate the arithmetic mean to measure the global similarity between two-frame objects. The closer to 1, the more similar they are.

[0068] The execution method of the key-frame motion detection method is as follows: Calculate the similarity between the memory object generated by the non-key frame and the memory object of the nearest key frame saved in the memory bank. If it is less than the threshold δ, the current frame is re-labeled as a key frame, a prompt is generated and the final mask is re-predicted, and then a new memory object is generated.

[0069] The execution method of the redundant memory object pruning method is as follows: After the number of memory objects in the memory bank reaches the preset upper limit, the similarity between the memory object and all objects in the memory sequence of the non-key nearest frame in the memory bank is calculated, and the two memory objects with the largest similarity and greater than the threshold ε will be pruned.

[0070] Compared with the prior art, the present invention preferably solves the problems that the SAM2 model requires high-quality prompts to generate high-quality segmentation results, and in scenarios where automation is required in practical applications, it is necessary to solve the problems of automatic prompt generation, when to give prompts, and the introduced additional computational consumption problems. The present invention annotates the screened zebrafish behavior images and uses the label images as the training data of the fully automatic segmentation model, without the need to annotate each frame of the video data, greatly reducing the cost of manual annotation. During the model training process, through the loss function under various conditions, the generated prompts are more accurate, and the image encoder and mask decoder can also better adapt to the new data after being fine-tuned using Adapter and LoRA respectively. The processing effect is better in the final predicted boundary details and artifact recognition. During the model inference process, through the self-identification of non-key frames, key frames are dynamically marked, avoiding unnecessary prompt generation for all frames, and the segmentation speed of the model is accelerated through dynamic pruning of redundant memories. The method is simple and effective, especially suitable for real-time segmentation and tracking of high-speed moving targets in microscope videos, and has good application prospects and commercial development value. Brief Description of the Drawings

[0071] Figure 1 It is a schematic diagram of the overall process of the embodiment of the present invention;

[0072] Figure 2 It is a schematic diagram of the structure of the fully automatic segmentation model;

[0073] Figure 3 It is a processing flow chart of the memory bank detection optimization module;

[0074] Figure 4 It is a comparison chart of the experimental results of Example 1 and other segmentation methods. Detailed Embodiments

[0075] The present invention is further described in detail below in conjunction with specific embodiments and drawings. Except for the contents specifically mentioned below, the processes, conditions, experimental methods, etc. for implementing the present invention are common knowledge and common common sense in the field and are not particularly limited in the present invention.

[0076] Example 1

[0077] See also Figure 1 , a fully automatic segmentation procedure for zebrafish behavior videos based on the SAM2 model is as follows:

[0078] S1: Obtain a collection of zebrafish behavior videos, screen a part of the frame image collection, annotate the zebrafish foreground on these frame images, and obtain a frame image label collection. The frame image collection is formed by manual screening, and mainly selects images that are difficult to segment due to artifacts caused by short-term intense exercise, target disappearance, and background impurity interference. Use labelme software to accurately annotate these images and mark the foreground area of ​​the zebrafish. Among them, artifacts caused by high-speed movement and background impurities covered by the translucent foreground are also considered to be foreground areas. Generate the corresponding segmentation mask, and finally review the annotation results to ensure the consistency and accuracy of the annotation. Among all the training data, 75% of the data will be divided into training sets, and the remaining 25% of the data will be divided into verification sets. Since the final task will be applied to the video data, the image data is not divided into a test set.

[0079] S2: Build a fully automatic segmentation model, train the model by screening the frame image set and the corresponding label set, and obtain the trained model parameters.

[0080] See also Figure 2 The overall structure of the model used in the present invention, some of which are from the SAM2 model, the SAM2 model includes: an image encoder, a prompt encoder, a mask decoder, a memory encoder, a memory bank and a memory attention module; the image encoder and the memory bank are both connected to the memory attention module; the memory attention and prompt encoder are both connected to the mask decoder; the prediction mask output by the mask decoder is connected to the memory bank through the memory encoder.

[0081] The fully automatic segmentation model constructed by the present invention introduces a prompt automatic generation module on the basis of the SAM2 model, including: learnable prompt tokens, a bidirectional attention module, a bounding box prediction multi-layer perceptron, and a rough mask prediction convolutional neural network; the learnable prompt tokens, the image feature embedding vectors generated by the image encoder, and their corresponding position embedding vectors are input into the bidirectional attention module, and part of the output content is input into the prompt encoder after predicting the bounding box and rough mask prompts through the multi-layer perceptron and the convolutional neural network respectively, and part of the output content is directly connected to the mask decoder as high-dimensional information. In addition, a memory bank detection optimization module is introduced on the basis of the SAM2 model, which includes: non-key frame motion detection and redundant memory object pruning. Finally, the image encoder of the SAM2 model is fine-tuned by adding an Adapter, and the mask decoder is fine-tuned by adding LoRA.

[0082] The names and functions of the various component parts in the fully automatic segmentation model are as follows:

[0083] 1), Image encoder: It consists of a hierarchical vision Transformer architecture Hiera, and uses a streaming processing method to extract features from all frame images, generate image feature embedding vectors, and is fine-tuned using an Adapter.

[0084] 2), Prompt encoder: Processes the prompts given manually and generated by the prompt automatic generation module, and generates prompt embedding vectors to be input into the mask decoder.

[0085] 3), Mask decoder: Receives the image feature embedding vectors and prompt embedding vectors, and generates the final segmentation mask, which is fine-tuned using LoRA.

[0086] 4), Memory encoder: Receives the final segmentation mask, fuses the image feature information, and generates memory objects.

[0087] 5), Memory bank: Stores all key frames and memory objects generated by the most recent k non-key frames.

[0088] 6), Memory attention module: Used to extract historical frame-related information from the memory bank for the current frame.

[0089] 7), Learnable prompt tokens: Used to learn and record the prompt information required for the target objects used by the model for automatic segmentation.

[0090] 8), Bidirectional attention module: Used for mutual information extraction and interaction from the image feature embedding vectors, image position embedding vectors, and learnable prompt tokens.

[0091] 9), Bounding box prediction multi-layer perceptron: Used to predict the bounding box prompts generated by the prompt automatic generation module.

[0092] 10) Coarse Mask Prediction Convolutional Neural Network: Used to predict the coarse mask prompts generated by the prompt automatic generation module.

[0093] 11) Non-critical Frame Motion Detection: By comparing the memory objects of non-critical frames with the objects of the nearest critical frame in the memory bank, determine whether the current frame needs to be regarded as a critical frame to generate prompts.

[0094] 12) Redundant Memory Object Pruning: By comparing the memory objects of the current frame with the objects of the nearest 6 frames in the memory bank, determine whether there is memory redundancy in the current memory bank.

[0095] The specific steps of the model training in step S2 are as follows:

[0096] S2-1: Input the frame image sample set into the full-automatic segmentation model, and obtain the prediction results of the bounding box and the coarse mask through the image encoder and the prompt automatic generation module in sequence. For the input samples of the training set and the corresponding frame image labels, due to the diverse actual scenarios, a series of strong data augmentation operations are adopted, including: random rotation of 90 degrees, random translation of 20%, random scaling and cropping with a magnification of 0.8 - 1.0, horizontal flipping with a probability of 30%, central cropping with a probability of 50%, randomly brightening / dimming the overall picture brightness by 50%, randomly adding Gaussian noise, a multi-transformation combination of random CLAHE operations, and adopting the form of dynamic adjustment plus short-side padding to maintain the size of the graphic training, and the size of the graphic training is 1024×1024.

[0097] The prompt automatic generation module will receive the multi-level image feature embedding vectors output by the image encoder, and perform information interaction through 3 bidirectional attention modules and learnable prompt vectors in sequence; the bidirectional attention module has 4 layers, the first layer is the self-attention layer, which is used for the learnable prompt token to capture its own context information; the second layer is the token-image cross-attention layer, which is used for the learnable prompt token to extract information from the image features; the third layer is the multi-layer perceptron layer, which performs non-linear transformation on the updated learnable prompt tokens through two fully connected layers; the fourth layer is the image-token cross-attention layer, which is used to inject the information of the learnable prompt tokens back into the image features.

[0098] The image encoder Hiera is hierarchically designed, so the features from low resolution to high resolution enter the bidirectional attention module for calculation in sequence, and are used as components of the query and key of the new layer. The obtained multi-level image feature embedding vectors are fused through the Feature Pyramid Network (FPN) to become multi-scale feature vectors, and a lightweight mask prompt prediction convolutional neural network is used to obtain the coarse mask prompt. The first two labels of the finally obtained updated prompt tokens are used to predict the upper left coordinate and the lower right coordinate of the bounding box prompt, and the bounding box prompt is obtained.

[0099] S2-2: Input the prediction results of the bounding box hint and the rough mask hint into the hint encoder, and obtain the final predicted mask through the mask decoder.

[0100] Specifically, the mask hint obtained in step S2-1 passes through the hint encoder to obtain a dense hint embedding vector; the upper left coordinate and the lower right coordinate of the predicted bounding box hint obtained in step S2-1 pass through the hint encoder and are concatenated with the remaining (T-2)×256-dimensional high-dimensional hint information to form a sparse hint embedding vector. After inputting the dense hint embedding vector and the sparse hint embedding vector into the mask decoder, the final predicted mask of the current frame is output.

[0101] S2-3: Calculate the generalized intersection over union loss through the predicted bounding box result and the ground truth bounding box generated by the frame image label. In addition to measuring the overlap degree between the predicted bounding box and the ground truth bounding box, the generalized intersection over union loss function (GIoU Loss) incorporates the influence of the area of the smallest convex region containing both, avoiding too obvious differences between the two.

[0102] The generalized intersection over union loss function L of the bounding box box is calculated by the following formula:

[0103]

[0104] where A and B are the areas of the ground truth bounding box and the predicted bounding box hint; C is the area of the smallest convex region containing A and B.

[0105] S2-4: Calculate the focal loss through the rough mask prediction result and the ground truth mask of the frame image. The focal loss function (Focal Loss) pays more attention to the task dominated by difficult samples than the traditional binary cross-entropy loss, guiding the segmentation of blurred regions. The focal loss L of the rough mask prediction result imask is calculated by the following formula:

[0106] L imask = L focal = -αy(1 - p) γ log(p) - (1 - α)(1 - y)p γ log(1 - p).

[0107] where y is the ground truth mask, p is the output result of the rough mask prediction through the sigmoid function, which is calculated pixel by pixel, and α, γ are hyperparameters. Preferably, α = 0.7 and γ = 3.

[0108] S2-5: Calculate the Dice loss, focal loss, and boundary loss through the final prediction mask and the frame image label ground truth mask. The Dice loss function is directly aligned with the evaluation metric and can better guide the final training effect of the model, L dice The calculation formula is:

[0109]

[0110] where P is the final prediction mask region and G is the ground truth mask region;

[0111] The calculation logic of the boundary loss function is to draw normal vectors for all points on the edge of the ground truth mask region and sum the L2 distances of the intersecting line segments with the edge of the prediction mask region. However, during the process of using it as a loss function, since the prediction mask region is indirectly represented by the softmax output and not explicitly defined, differential operations cannot be performed on the boundary points. After mathematical derivation, it is transformed into the integral of the shortest signed distance from any point in the difference region to the edge of the ground truth mask, L bound The calculation formula is:

[0112] L bound = ∫ Ω φ G (q)s(q)dq.

[0113] where s(q) represents the predicted softmax value in the image region, is the shortest signed distance from point q to the edge of the ground truth prediction.

[0114] The formula for the final prediction mask loss L omask is:

[0115] L omask = λ focal L focal + λ dice L dice + λ bound L bound .

[0116] where the calculation method of L focal is the same as the focal loss of the rough mask prediction result, y is the ground truth mask, p is the final prediction mask of the model, and all λ are hyperparameters. Preferably, λ focal = 10, λ dice = 1, λ bound = 1.

[0117] The formula for the mask loss Lmask that combines the final prediction mask loss and the rough mask prediction loss is:

[0118] L mask = λ i Limask +λ o L omask 。

[0119] Among them, all λ are hyperparameters. Preferably, λ i = 10, λ o = 1.

[0120] S2-6: Calculate the cross-entropy loss by predicting whether the target object is contained in the decoder and whether the target object is contained in the true mask.

[0121] In order to determine whether the current frame image has a target object, the mask decoder of the SAM2 model will output a target score. Calculate the cross-entropy loss based on whether the score and the true label contain the target. The loss L obj for whether the target object is contained is calculated as follows:

[0122] L obj = -ylog(p) - (1 - y)log(1 - p).

[0123] Among them, y is the occurrence of the target object in the true mask, and p is the sigmoid value of the target object score output by the mask decoder.

[0124] S2-7: Calculate the total loss by weighted calculation of all losses in steps S2-3 to S2-6. The calculation formula of the total loss function is:

[0125] L total = λ1L box + λ2L mask + λ3L obj 。

[0126] Among them, all λ are hyperparameters. Preferably, λ1 = λ2 = λ3 = 1.

[0127] S2-8: The AdamW optimizer performs backpropagation on the fully automatic segmentation model through the total loss, and updates the parameters of the prompt automatic generation module in the model and the newly added parameters of the mask decoder LoRA.

[0128] Specifically, by combining the calculations of multiple loss functions, judge the prediction effects of the model on the bounding box prompt prediction, rough mask prompt prediction, final mask prediction, and target appearance detection. After obtaining the loss function, the optimizer can be used to calculate backpropagation and update the weights. Calculate the Dice coefficient between the final prediction result and the true label on the validation set in each round, and save the model parameters with the highest Dice coefficient. Preferably, the initial learning rate of the AdamW optimizer is set to 0.0001, and the learning rate is adjusted using cosine annealing, with a minimum value of 0.00001.

[0129] S2-9: Repeat steps S21 - S28 up to 200 rounds, and select the one with the highest true mask Dice coefficient on the validation set as the optimal model parameters for the subsequent process.

[0130] S3: Obtain the zebrafish behavior video to be segmented. After converting it into a frame image sequence using the FFmpeg program, input it into the fully automatic segmentation model that loads the trained model parameters. Utilize the memory mechanism to propagate historical frame information, and determine the timing of automatically generating prompts and reduce redundant operations through the memory bank detection and optimization module to obtain the final segmentation result.

[0131] The specific steps of using the trained model in step S3 for the zebrafish behavior video segmentation task are as follows:

[0132] S3-1: Convert the zebrafish behavior video data into a frame image sequence, input it into the trained fully automatic segmentation model, obtain all frame image feature embedding vectors through the image encoder, and mark the first frame as the key frame, and the rest as non-key frames.

[0133] Specifically, place the frame image sequence of the zebrafish behavior video in the same parent folder, name it according to a 4-digit number with leading zeros, and the file suffix is png or jpg. This method will read all the images in the file path and sort them. After the model loads the trained parameters, all frame images are sequentially passed through the image encoder by means of streaming processing to obtain the image feature embedding vectors.

[0134] S3-2: If the current frame is a non-key frame, the frame image feature embedding vector fuses the memory bank information through memory attention and inputs it into the mask decoder to obtain the predicted mask; if the current frame is a key frame, the frame image feature embedding vector will also generate a prompt embedding vector through the prompt automatic generation module, and together with the image features containing memory information, input it into the mask decoder to obtain the final predicted mask.

[0135] Specifically, the frame image feature embedding vector performs cross-attention calculation with the objects in the memory bank to obtain the feature embedding vector containing temporal information, and obtains the predicted mask through the mask decoder.

[0136] Furthermore, the first frame is the first key frame. The feature embedding vector of the corresponding frame is encoded by position encoding to obtain the position embedding vector. The position embedding vector and the feature embedding vector are added element-wise as the key, and the learnable prompt token is used as the query. Information is fused through the bidirectional attention module to update the learnable prompt token. The learnable prompt token is a vector of T×256, where T is the total number of labels used for object recognition in the token, which is set to 10 here. In addition, if it is a key frame other than the first frame, the position embedding vector also performs cross-attention calculation with the objects in the memory bank to fuse temporal information. Correspondingly, since there is no temporal information for reference in the first frame, zero vectors of the corresponding size are used for padding.

[0137] S3-3: For the predicted masks and original image features of all frames, generate memory objects through the memory encoder.

[0138] See Figure 3 , specifically, each memory object is obtained by downsampling the predicted mask, directly adding it to the image features projected by a 1×1 convolution, and then fusing multi-level features through multiple stacked depthwise separable convolutional layers. Therefore, each memory object in the memory bank records the image features of the corresponding frame and the final predicted mask information of the corresponding frame, which plays an important role in the subsequent propagation of memory information.

[0139] S3-4: If the current frame is a non-key frame, execute the non-key frame motion detection method. If it is relabeled as a key frame, go back to step S3-2 to execute again. Otherwise, save the memory object to the memory sequence of the nearest non-key frame, and slide the memory window; if the current frame is a key frame, save the memory object to the key frame memory sequence.

[0140] The main purpose of this step is to avoid the additional computational consumption caused by generating prompts for each frame. Therefore, the cosine similarity is used to quantify the dynamic change information between the current frame and the nearest key frame.

[0141] Specifically, the size of the memory object is B×C×H×W. During the inference process, B = 1, C = H = W = 64. When calculating the similarity, the spatial dimension of the object is flattened, and the cosine similarity is calculated position-wise along the channel dimension in space, and finally the arithmetic mean is taken to measure the global similarity between two frames. The larger the value, the more similar the overall. The similarity calculation formula is:

[0142]

[0143] where Ai and Bi are the channel vectors at the corresponding positions after the memory object is flattened in the spatial dimension.

[0144] Furthermore, since the current frame memory object integrates the image features of the current frame and the information of the final prediction mask of the current frame, comparing the similarity with the memory object of the nearest key frame helps to determine whether there are significant changes in the current frame compared to the previous key frame and the differences in mask prediction. If the differences are obvious, the current frame is marked as a key frame, and a hint is generated for the current frame by the hint automatic generation module. After obtaining a new mask and generating a new memory object, it is added to the memory bank. The update method of the memory bank is as shown from time t+1 to time t+2 in Figure 3 to determine the selection of the threshold through the visual results of the experiment. Preferably, the threshold δ = 0.9 here.

[0145] S3-5: Execute the redundant memory object pruning method, and the pruned memory object will not be used for the memory attention operation of the next frame.

[0146] The main purpose of this step is to avoid a large amount of redundant information in the nearest k-frame memory objects of the memory bank when the zebrafish is stationary for a long time, resulting in additional computational resource requirements and affecting the computational efficiency.

[0147] Furthermore, the redundant memory module is implemented before the next frame passes through the memory attention module, that is, the similarity between the nearest frame in the memory bank and the memory objects of the previous at most k-1 frames is calculated. The calculation method of cosine similarity is the same as that in step S34, and the spatial global overall similarity of two memory objects is obtained. To avoid the misdeletion of key information caused by continuous movement within a period of time, it is designed that only the memory objects with a similarity greater than the threshold ε will be removed. The update method of the memory bank is as shown from time t to time t+1 in Figure 3 to determine the selection of the threshold through the visual results of the experiment. Preferably, the threshold ε = 0.95 here.

[0148] S3-6: Perform the operations of S3-2 to S3-5 on each frame of the video. After all frames are processed, the final zebrafish behavior video segmentation result is obtained.

[0149] To verify the effectiveness and accuracy of this embodiment, a typical motion scenario was selected, and other segmentation methods were selected for experimental comparison with this embodiment.

[0150] Refer to Figure 4 , compared with the decoupled video object segmentation scheme DEVA and SAM2Rad without introducing a memory mechanism, the present invention can obtain more accurate segmentation results closer to the real labels when facing complex motions such as motion tail flicks.

[0151] The above specific implementation is only a further description of the present invention, and is not intended to limit the patent of the present invention. All equivalent implementations of the present invention should be included within the scope of the claims of the present invention.

Claims

1. A fully automatic segmentation method for zebrafish behavior videos based on the SAM2 model, comprising the SAM2 model and its image encoder and mask decoder, characterized in that: The SAM2 model introduces a prompt automatic generation module and a memory library detection optimization module, and adds Adapter to the image encoder for optimization, and adds LoRA to the mask decoder for optimization. The module inputs the generated prompt information into the memory library optimization module to obtain the fully automatic segmentation result of the zebrafish behavior video. The specific steps include: S1: Obtaining the zebrafish behavior image dataset Select a set of frames with strong motion features and impurity interference from the zebrafish behavior video, annotate the zebrafish foreground in these frames, and also annotate the artifacts caused by high-speed motion. S2: Build and train a fully automatic segmentation model The model is trained by screening the frame image set and the corresponding label set, and the trained model parameters are used to construct a fully automatic segmentation model; S3: Using the model to segment zebrafish behavior videos The zebrafish behavior video is converted into a frame image sequence using the FFmpeg program and then input into the fully automatic segmentation model. The memory mechanism is used to propagate historical frame information to obtain the segmentation result.

2. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 1, characterized in that: The automatic prompt generation module includes: a learnable prompt tag, a bidirectional attention module, a bounding box prediction multi-layer perceptron, and a rough mask prediction convolutional neural network; the image feature embedding vector generated by the learnable prompt tag and the image encoder and its corresponding position embedding vector are input into the bidirectional attention module, and its output is passed through the multi-layer perceptron and the convolutional neural network to obtain the predicted bounding box and rough mask prompt, which are input into the prompt encoder, and the output of the prompt encoder is connected to the mask decoder as high-dimensional information.

3. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 1, characterized in that: The memory library detection optimization module includes: non-key frame motion detection and redundant memory object pruning.

4. The method for fully automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 1, characterized in that: The step S2 specifically includes: S2-1: Input the frame image sample set into the fully automatic segmentation model, and pass through the image encoder and the prompt automatic generation module to obtain the prediction results of the bounding box and the rough mask; S2-2: The prediction results of the bounding box and the rough mask are input into the hint encoder, and the final prediction mask is obtained after passing through the mask decoder; S2-3: Calculate the generalized intersection-over-union loss using the bounding box prediction result and the ground-truth bounding box generated by the frame image label; S2-4: Calculate the focal loss using the coarse mask prediction result and the true mask of the frame image label; S2-5: Calculate Dice loss, focal loss and boundary loss through the final predicted mask and the real mask of the frame image label; S2-6: Calculate the cross entropy loss by predicting whether the decoder contains the target object and whether the real mask contains the target object; S2-7: Calculate the total loss by weighted calculation of all losses in steps S2-3 to S2-6; S2-8: The AdamW optimizer back-propagates the fully automatic segmentation model through the total loss and updates all parameters in the model that are not frozen; S29: Repeat steps S2-1 to S2-8 multiple times, and select the one with the highest true mask Dice coefficient on the validation set as the optimal model parameter, wherein the multiple repetitions are less than 200 times.

5. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 4 is characterized in that: The generalized intersection-over-union loss function L of the bounding box in step S2-3 box Calculated by the following formula: Among them, A and B are the areas of the real bounding box and the predicted bounding box prompt; C is the area of ​​the smallest convex region containing A and B; The focal loss L of the rough mask prediction result in step S2-4 imask Calculated by the following formula: L imask =L focal =-αy(1-p) γ log(p)-(1-α)(1-y)p γ log(1-p); Among them, y is the real mask; p is the output result of the rough mask prediction through the sigmoid function; α and γ are hyper parameters; The final prediction mask loss L of step S2-5 is omask Calculated by the following formula: L omask =λ focal L focal +λ dice L dice +λ bound L bound ; Among them, L focal The focal loss is calculated in the same way as the coarse mask prediction result; y is the real mask; p is the final predicted mask of the model; λ is a hyperparameter; L dice is the Dice loss function; L bound is the boundary loss function; The Dice loss function L dice Calculated by the following formula: Among them, P is the final predicted mask area; G is the real mask area; The boundary loss function L bound To construct the normal line for all points on the edge of the true mask area, the sum of the L2 distances of the line segments intersecting the edge of the predicted mask area can be transformed into the integral of the shortest signed distance from any point in the difference area to the true mask edge: L bound =∫Ωφ G (q)s(q)dq; Among them, s(q) is the predicted softmax value in the image area; is the shortest signed distance from point q to the true predicted edge.

6. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 5, characterized in that: The mask loss L of the final prediction mask loss and the rough mask prediction loss mask Calculated by the following formula: L mask =λ i L imask +λ o L omask ; Among them, λ is a hyperparameter.

7. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 4, characterized in that: Whether the step S2-6 contains the loss L of the target object obj Calculated by the following formula: L obj =-ylog(p)-(1-y)log(1-p); Among them, y is the occurrence of the target object in the real mask; p is the sigmoid value of the target object score output by the mask decoder.

8. The method for automatic segmentation of zebrafish behavior videos based on the SAM2 model according to claim 4, characterized in that: The total loss function of step S2-7 is calculated by the following formula: L total =λ1L box +λ2L mask +λ3L obj ; Here, all λ are hyperparameters.

9. The method for fully automatic segmentation of zebrafish behavior videos based on SAM2 according to claim 1, characterized in that: The step S3 is specifically as follows: S3-1: Convert zebrafish behavior video data into frame image sequences, input the trained fully automatic segmentation model, obtain all frame image feature embedding vectors through the image encoder, and mark the first frame as a key frame, and the rest are non-key frames; S3-2: If the current frame is a non-key frame, the frame image feature embedding vector fuses the memory library information through memory attention and is input into the mask decoder to obtain the predicted mask; if the current frame is a key frame, the frame image feature embedding vector will also generate a hint embedding vector through the hint automatic generation module, and input it into the mask decoder together with the image features containing memory information to obtain the final predicted mask. S3-3: For the predicted masks and original image features of all frames, a memory object is generated through the memory encoder; S3-4: If the current frame is a non-key frame, execute the non-key frame motion detection method. If it is re-marked as a key frame, return to step S3-2 for execution. Otherwise, save the memory object to the most recent non-key frame memory sequence and slide the memory window. If the current frame is a key frame, save the memory object to the key frame memory sequence. S3-5: Execute the redundant memory object pruning method. The pruned memory objects will not be used for the memory attention operation of the next frame. S3-6: Execute the operations of S3-2 to S3-5 for each frame of the video. After all frames are processed, the final zebrafish behavior video segmentation result is obtained.

Images