Geometric primitive detection model training method and device based on key points

Abstract

The invention provides a geometric primitive detection model training method and device based on key points, and the method comprises the steps: determining the multi-scale image features of a sample image based on an initial detection model, carrying out the fusion of the multi-scale image features, carrying out the geometric primitive detection based on the instance query features obtained through fusion, and obtaining a geometric primitive detection model. Obtaining a detection result of each prediction geometric primitive in the sample image; and performing parameter iteration on the initial detection model based on the detection result of each prediction geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model, thereby overcoming the defect of poor generalization performance of a geometric primitive detection method in a traditional scheme, building the initial detection model through an end-to-end detection framework, and improving the detection accuracy of the geometric primitive. The problems that the geometric primitives are difficult to accurately express and the detection precision is low in a traditional target detection scheme can be avoided, the generalization performance of the network model can be enhanced, and the stability and the detection precision of the network model can be improved.

Description

Technical Field

[0001] The present invention relates to the field of computer vision technology, and in particular to a key point-based geometric primitive detection model training method and device. Background Art

[0002] Geometric primitives frequently appear in flowcharts, each representing a specific meaning. By combining simple geometric primitives, we can construct semantically complex flowcharts. Therefore, geometric primitive detection plays a crucial role in advanced vision tasks such as automated document processing, flowchart recognition, and semantic understanding.

[0003] Currently, geometric primitive detection relies heavily on manually defined rules or hand-extracted features. While these methods can perform recognition-based detection, they suffer from poor generalization performance and can lead to false detections and missed detections in complex scenarios. With the development of deep learning technology, neural network-based geometric primitive detection methods have been proposed. While these methods can detect the location of primitives, they struggle to precisely determine the geometric primitives and their parameters. Summary of the Invention

[0004] The present invention provides a method and device for training a geometric primitive detection model based on key points, which is used to solve the defects of low accuracy of geometric primitive detection and difficulty in accurately representing geometric primitives in the prior art. An initial detection model is built using an end-to-end detection framework. By training the initial detection model, a geometric primitive detection model with superior performance in geometric primitive detection is obtained, thereby realizing high-precision geometric primitive detection based on the geometric primitive detection model.

[0005] The present invention provides a key point-based geometric primitive detection model training method, comprising:

[0006] Determine a sample image, wherein the sample image includes a plurality of sample geometric primitives;

[0007] Determining multi-scale image features of the sample image based on the initial detection model, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results of each predicted geometric primitive in the sample image;

[0008] Based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, the initial detection model is iterated on parameters to obtain a geometric primitive detection model.

[0009] According to a key point-based geometric primitive detection model training method provided by the present invention, the initial detection model includes a feature extraction network, a feature fusion network, and a primitive detection network;

[0010] The feature extraction network is used to perform multi-scale feature extraction on the sample image to obtain multi-scale image features of the sample image;

[0011] The feature fusion network is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features;

[0012] The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain detection results of each predicted geometric primitive in the sample image.

[0013] According to a key point-based geometric primitive detection model training method provided by the present invention, the feature fusion network includes an encoder and a decoder;

[0014] The encoder is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features to obtain a fused feature;

[0015] The decoder is used to perform cross-attention calculation based on the fused features and an initial query vector obtained by random initialization to obtain instance query features.

[0016] According to a keypoint-based geometric primitive detection model training method provided by the present invention, the detection results include the category of the corresponding predicted geometric primitive, the center point, the key point, the angle and distance of the key point relative to the center point, and the corresponding relationship between each sector area and the key point;

[0017] The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain the category and center point of each predicted geometric primitive and the corresponding relationship between each sector area and the key point;

[0018] The primitive detection network is further configured to determine, when a key point of any predicted geometric primitive exists in any sector-shaped area corresponding to any predicted geometric primitive, an angle and a distance of the key point of the predicted geometric primitive relative to the center point, and determine the key point of the predicted geometric primitive based on the angle and the distance of the key point of the predicted geometric primitive relative to the center point;

[0019] Among them, any prediction geometric primitive corresponds to multiple fan-shaped areas, and the central angles of the fan-shaped areas are the same.

[0020] According to a keypoint-based geometric primitive detection model training method provided by the present invention, the sample labels include the category label of the corresponding sample geometric primitive, the center point, the key point, the angle and distance of the key point relative to the center point, and the relationship label between each corresponding sector area and the key point; any sample geometric primitive corresponds to multiple sector areas, and the central angle of each sector area is the same;

[0021] The method of performing parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model includes:

[0022] Determining a classification loss based on the category of each predicted geometric primitive, the correspondence between each sector area corresponding to each predicted geometric primitive and the key point, and the category label of each sample geometric primitive and the relationship label between each sector area corresponding to each sample geometric primitive and the key point;

[0023] Determining a regression loss based on the center point of each predicted geometric primitive, the angle and distance of the key point relative to the center point, and the angle and distance of the center point of each sample geometric primitive, the angle and distance of the key point relative to the center point;

[0024] Determining a coincidence loss based on the center points and key points of each predicted geometric primitive and the center points and key points of each sample geometric primitive;

[0025] Based on the classification loss, the regression loss and the coincidence loss, parameter iteration is performed on the initial detection model to obtain a geometric primitive detection model.

[0026] According to a keypoint-based geometric primitive detection model training method provided by the present invention, determining the coincidence loss based on the center point and keypoint of each predicted geometric primitive and the center point and keypoint of each sample geometric primitive includes:

[0027] Determining a degree of overlap between each predicted geometric primitive and the corresponding sample geometric primitive based on a distance between a center point of each sample geometric primitive and a key point, a distance between a center point of each sample geometric primitive and a key point of a corresponding predicted geometric primitive, a distance between a center point of each predicted geometric primitive and a key point, and a distance between a center point of each predicted geometric primitive and a key point of a corresponding sample geometric primitive;

[0028] Based on the overlap degrees corresponding to the predicted geometric primitives, an overlap loss is determined.

[0029] According to a key point-based geometric primitive detection model training method provided by the present invention, the sample image is determined based on the following steps:

[0030] Acquire an initial image set, wherein the initial image set includes a plurality of initial images;

[0031] Preprocessing is performed on each initial image in the initial image set to obtain a sample image set, where the sample image set contains multiple sample images; the preprocessing includes at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization.

[0032] The present invention also provides a key point-based geometric primitive detection model training device, comprising:

[0033] a determining unit, configured to determine a sample image, wherein the sample image includes a plurality of sample geometric primitives;

[0034] a detection unit, configured to determine multi-scale image features of the sample image based on an initial detection model, fuse the multi-scale image features, and perform geometric primitive detection based on the fused instance query features to obtain a detection result of each predicted geometric primitive in the sample image;

[0035] The iterative unit is used to perform parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model.

[0036] The present invention also provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and runnable on the processor, wherein when the processor executes the program, it implements any of the key point-based geometric primitive detection model training methods described above.

[0037] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements any of the key point-based geometric primitive detection model training methods described above.

[0038] The key point-based geometric primitive detection model training method and device provided by the present invention determine the multi-scale image features of the sample image through the initial detection model, fuse the multi-scale image features, and perform geometric primitive detection based on the instance query features obtained by the fusion to obtain the detection results of each predicted geometric primitive in the sample image; based on the detection results of each predicted geometric primitive and the sample labels of each sample geometric primitive, a geometric primitive detection model is trained to obtain automatic geometric primitive detection based on the geometric primitive detection model, overcoming the defect of poor generalization performance of the geometric primitive detection method in the traditional scheme. The initial detection model is built through the end-to-end detection framework, and the geometric primitive detection model is obtained by training. This not only avoids the problems of difficult accurate representation of geometric primitives and low detection accuracy in traditional target detection schemes, but also enhances the generalization performance of the network model and improves its stability and detection accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0039] In order to more clearly illustrate the technical solutions in the present invention or the prior art, a brief introduction will be given below to the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0040] Figure 1 Schematic diagram of the process of the geometric primitive detection model training method based on key points provided by the present invention;

[0041] Figure 2 This is a flowchart of a geometric primitive detection model training method based on key points provided by the present invention;

[0042] Figure 3 is an example diagram of the geometric primitive detection process provided by the present invention;

[0043] Figure 4 This is an example diagram of the calculation of the degree of coincidence between the predicted geometric primitive and the sample geometric primitive provided by the present invention;

[0044] Figure 5 This is the overall framework diagram of the key point-based geometric primitive detection model training method provided by the present invention;

[0045] Figure 6 Schematic diagram of the structure of the key point-based geometric primitive detection model training device provided by the present invention;

[0046] Figure 7 It is a structural schematic diagram of the electronic device provided by the present invention. DETAILED DESCRIPTION

[0047] To make the objectives, technical solutions, and advantages of the present invention more clear, the technical solutions of the present invention will be clearly and completely described below in conjunction with the accompanying drawings. Obviously, the embodiments described are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts shall fall within the scope of protection of the present invention.

[0048] Early geometric primitive detection usually relied on manually set rules or achieved primitive recognition by extracting manual features. Such methods had poor generalization performance and could cause problems such as false detection and missed detection in some complex scenarios.

[0049] With the development of deep learning technology, researchers have used convolutional neural networks and attention mechanisms to automatically extract features of objects in images, enabling adaptive detection of their locations. However, traditional object detection methods typically use a rectangular box to outline the object in the image. However, for geometric primitives with multiple vertices, such as triangles and pentagons, a rectangular box cannot accurately represent their shape and location.

[0050] To this end, the present invention provides a key point-based geometric primitive detection model training method, which aims to train a geometric primitive detection model on the basis of an end-to-end detection framework and ensure the superior performance of the model in geometric primitive detection, thereby achieving high-precision geometric primitive detection. Figure 1 : is a flow chart of the geometric primitive detection model training method based on key points provided by the present invention, such as Figure 1 As shown, the method includes:

[0051] Step 110, determining a sample image, wherein the sample image includes a plurality of sample geometric primitives;

[0052] Specifically, before model training, the data required for model training, namely the training dataset, must be determined. Since geometric primitives are components of flow charts, current geometric primitive recognition often uses flow charts as a carrier. By identifying these primitives, information such as the geometric primitives in the image and their categories can be obtained. Based on this, in embodiments of the present invention, the training dataset required for model training can be determined as an image set, namely a sample image set.

[0053] Here, the sample image set contains multiple sample images, and each sample image contains multiple sample geometric primitives. The types and quantities of the sample geometric primitives are not exactly the same, that is, each sample image can contain sample geometric primitives of different categories, and the number of sample geometric primitives of each category can be the same or different. This can ensure the difference and diversity of the distribution of sample geometric primitives in the sample image set, thereby realizing the diversification of model training data, and then in the subsequent training process, the generalization ability of the model can be enhanced and the performance of the model can be improved.

[0054] The categories of the sample geometric primitives can be rectangle, triangle, rhombus, circle, ellipse, pentagon, hexagon, arrow, etc. In a sample image containing a complete flowchart, each sample geometric primitive is randomly connected to another sample geometric primitive through an arrow.

[0055] It is worth noting that after determining the image set, it is also necessary to annotate it. For example, the category, center point, key point, distance and angle between key points and center point of each sample geometric primitive in each sample image in the sample image set, and whether there are key points in each corresponding sector area, etc., are annotated to form sample labels, which facilitate subsequent comparison and calculation with the detection results of the model. Here, the sample label can include the category label, center point, key point, angle and distance of the key point relative to the center point of the sample geometric primitive, and the relationship label between each corresponding sector area and the key point.

[0056] Key points can be polygon vertices, arrow starting points, and turning points, and other points that can represent the shape, size, and structure of a geometric primitive. Sector regions are formed by dividing the 360° range of the corresponding geometric primitive into multiple sectors, with the center point of the corresponding geometric primitive as the center. Each sector region has the same central angle.

[0057] Step 120: determining multi-scale image features of the sample image based on the initial detection model, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results for each predicted geometric primitive in the sample image;

[0058] Step 130 : Based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, the initial detection model is iterated to obtain a geometric primitive detection model.

[0059] Specifically, after obtaining the sample images for model training, step 120 and step 130 may be executed to perform model training accordingly to obtain a trained geometric primitive detection model.

[0060] It is understandable that after obtaining the sample image and before model training, the first thing to determine is the training object, that is, the initial model of the geometric primitive detection model obtained through training, which is referred to herein as the initial detection model. Considering that the goal of model training is to improve performance and achieve high-precision geometric primitive detection, the embodiments of the present invention can build an initial detection model based on the end-to-end detection framework, and use it to train a geometric primitive detection model. Based on this model, automated and high-precision geometric primitive detection can be achieved, avoiding the problems of manual intervention and post-processing in traditional solutions, making it simpler, more reliable, and more effective.

[0061] After determining the sample image and the initial detection model, in an embodiment of the present invention, the initial detection model can be first applied to extract effective information of geometric primitives in the sample image, thereby obtaining multi-scale image features of the sample image. That is, the initial detection model can be used to perform multi-scale feature extraction on the sample image to obtain multi-scale geometric primitive related information hidden in the sample image, thereby extracting multi-scale, multi-granularity and multi-level features, namely, multi-scale image features.

[0062] Next, the initial detection model is applied to fuse these multi-scale image features to obtain features containing both local and global contextual information, namely, instance query features. Specifically, an attention mechanism can be introduced during the feature fusion process to mimic the way the human eye observes an object. This mechanism automatically focuses on relevant and important information in the multi-scale image features and ignores irrelevant information during the feature fusion process to model global information. This results in semantically rich fused features containing both local and global contextual information, namely, instance query features.

[0063] Afterwards, based on this instance query feature, geometric primitive detection can be performed through the initial detection model to detect the various geometric primitives present in the sample image, as well as the relevant parameters of each geometric primitive, such as the center point, the angle and distance of the key point relative to the center point, etc., thereby obtaining the detection results of each geometric primitive in the sample image output by the model prediction, that is, the detection results of each predicted geometric primitive. Here, the detection result can be the category of the predicted geometric primitive, the center point, the corresponding relationship between each sector area and the key point, the angle and distance of the key point relative to the center point, the key point, etc. The definition of each sector area corresponding to the predicted geometric primitive is basically the same as the definition of each sector area corresponding to the sample geometric primitive above.

[0064] Furthermore, after the initial detection model completes the detection of the sample image, supervised learning can be performed based on the detection results output by the model and the sample labels. The parameters of the model are continuously adjusted during the training process so that the detection results output by the model are as close as possible to the sample labels, and finally a trained detection model can be obtained.

[0065] Specifically, here we can use the detection results of each predicted geometric primitive output by the initial detection model and the sample label of the sample image to determine the loss of the initial detection model in the entire process of geometric primitive detection on the sample image. That is, the various parameters in the detection results can be compared with the parameters in the sample label in the same dimension to judge the loss of the model in this geometric primitive detection task through the similarity or difference between the corresponding data.

[0066] Afterwards, the parameters of the initial detection model can be iterated based on this loss to adjust the parameters of the initial detection model so that when the adjusted initial detection model is facing the sample image, the detection results output by the sample image can be as consistent as possible with the sample label of the sample image, thereby reducing its loss in the detection task and improving the performance of the model in geometric primitive detection. The above process is repeated through multiple rounds of parameter iterations, and finally a geometric primitive detection model that performs extremely well in the geometric primitive detection task can be obtained.

[0067] After the training is complete, a trained geometric primitive detection model is obtained. The testing process can now be executed, calling this geometric primitive detection model to perform geometric primitive detection, thereby obtaining detection results for each geometric primitive in each test image in the test dataset. The detection results here can include the category, key points, etc. of each geometric primitive in the test image.

[0068] It is noteworthy that in the embodiments of the present invention, the distribution of the training dataset is taken into account when determining the test dataset. This ensures that the data distribution of the training and test datasets is identical. This allows the model's training effectiveness and performance in geometric primitive detection to be verified in the post-training test phase using the same test dataset. Specifically, a large geometric primitive dataset can be constructed first, and then divided into training and test datasets according to a predetermined ratio (e.g., 8:2, 7:3, etc.), ensuring that the training and test datasets follow the same data distribution.

[0069] The key point-based geometric primitive detection model training method provided by the present invention determines the multi-scale image features of the sample image through the initial detection model, fuses the multi-scale image features, and performs geometric primitive detection based on the instance query features obtained by the fusion to obtain the detection results of each predicted geometric primitive in the sample image; based on the detection results of each predicted geometric primitive and the sample labels of each sample geometric primitive, a geometric primitive detection model is trained to obtain automatic geometric primitive detection based on the geometric primitive detection model, overcoming the defect of poor generalization performance of the geometric primitive detection method in the traditional scheme. The initial detection model is built through the end-to-end detection framework, and the geometric primitive detection model is obtained by training. This not only avoids the problems of difficult accurate representation of geometric primitives and low detection accuracy in traditional target detection schemes, but also enhances the generalization performance of the network model and improves its stability and detection accuracy.

[0070] Based on the above embodiment, the initial detection model includes a feature extraction network, a feature fusion network, and a primitive detection network;

[0071] The feature extraction network is used to extract multi-scale features of the sample image to obtain multi-scale image features of the sample image;

[0072] The feature fusion network is used to fuse multi-scale image features based on the correlation between each image feature in the multi-scale image features;

[0073] The primitive detection network is used to perform geometric primitive detection based on instance query features and obtain detection results of each predicted geometric primitive in the sample image.

[0074] Specifically, in embodiments of the present invention, when constructing an initial detection model, a single-stage end-to-end detection framework can be employed, integrating a feature extraction network, a feature fusion network, and a geometric primitive detection network to construct the initial detection model. That is, the initial detection model can include a feature extraction network, a feature fusion network, and a primitive detection network.

[0075] Figure 2 This is a flow chart of the geometric primitive detection model training method based on key points provided by the present invention, such as Figure 2 As shown in the figure, after the sample image is input into the initial detection model, it first passes through the feature extraction network to perform multi-scale feature extraction to obtain effective multi-scale information, thereby obtaining multi-scale image features. Considering that the deep residual neural network can effectively capture the multi-scale information of the target in the image, in this embodiment of the present invention, a feature extraction network can be constructed based on the deep residual neural network ResNet.

[0076] Specifically, in an embodiment of the present invention, a series of convolutional layers are used to extract features from the input sample image. These convolutional layers use convolution kernels with different receptive fields and step sizes to adaptively extract effective information from the sample image, thereby obtaining multi-scale, multi-granularity and multi-level features. In this way, the initial detection model can adapt to geometric primitives of various deformations and postures, thereby providing data support for subsequent geometric primitive detection, thereby improving the detection performance of the model.

[0077] In addition, during the multi-scale feature extraction process, due to the addition of residual connections, the model can effectively perform gradient backpropagation, thereby achieving effective learning.

[0078] After the multi-scale image features are extracted by the feature extraction network, they are then passed through the feature fusion network using the attention mechanism to fuse them, resulting in richer and more discriminative features, namely instance query features. Specifically, the feature fusion network fuses the multi-scale image features based on the correlation between each of the image features within the multi-scale image features to produce instance query features.

[0079] Also, see Figure 2It can be seen that in the process of feature fusion, the attention mechanism is introduced in the embodiment of the present invention to construct a feature fusion network of the attention-based encoding and decoding architecture, namely, Transformer Encoder and Transformer Decoder.

[0080] Afterwards, the instance query features fused by the feature fusion network are fed into the primitive detection network for detection, thereby obtaining the detection results of each predicted geometric primitive in the sample image. Specifically, a keypoint-based geometric primitive detection network is constructed in polar coordinates. The instance query features containing global semantic information obtained by the previous fusion step are input into a feedforward neural network (FFN), and the detection results are obtained through a classifier and regressor.

[0081] Based on the above embodiment, the feature fusion network includes an encoder and a decoder;

[0082] The encoder is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features to obtain a fused feature;

[0083] The decoder is used to perform cross-attention calculation based on the fused features and the initial query vector obtained by random initialization to obtain instance query features.

[0084] Specifically, in the embodiment of the present invention, the feature fusion network includes an encoder Transformer Encoder and a decoder Transformer Decoder.

[0085] The encoder takes multi-scale image features extracted by the feature extraction network as input. The input multi-scale image features are of size C × HW, where C is the number of channels, and H and W are the height and width, respectively. After the multi-scale image features are input to the encoder, the encoder first performs a linear transformation on each image feature to obtain the three elements used to calculate the attention relationship: query, key, and value. Then, the correlation between each image feature and other image features is calculated according to the attention calculation formula.

[0086] Here, the attention calculation process can be expressed using the following formula:

[0087]

[0088] Where Q, K and V are query, key and value respectively, d k The dimension of the key, used to avoid QK T The product is too large.

[0089] When calculating the correlation between image features, Q comes from the current image feature, K and V come from other image features, QK T Indicates the similarity between the current image feature and other image features. Each image feature is calculated with other image features to obtain a similarity score vector. Finally, the similarity score vector is normalized using softmax to obtain the attention weight vector Attention(Q,K,V). This attention weight vector represents the degree of correlation between the current image feature and other image features.

[0090] Afterwards, the multi-scale image features can be fused according to the correlation degree to obtain fused features. Here, feature fusion based on the correlation degree can obtain richer and more discriminative features, namely fused features.

[0091] In an embodiment of the present invention, the attention mechanism can enable each image feature to interact with other image features, thereby obtaining semantically rich global context information. This process will be repeated multiple times, so that more and deeper global context information can be obtained.

[0092] Correspondingly, the decoder can use the fused features output by the encoder as input and decode the fused features to obtain the instance query features. That is, it can use the initial query vector Object Query obtained by random initialization and the fused features output by the encoder with global context information to perform cross-attention calculations to perceive the positions of each geometric primitive in the sample image. Here, the number of Object Queries corresponds to the number of geometric primitives in the sample image; each Object Query gradually adjusts its own representation through multiple layers of cross-attention and self-attention to match the features of a certain geometric primitive. Ultimately, each Object Query can perceive the category and location information of the geometric primitive, thereby facilitating subsequent geometric primitive detection.

[0093] The calculation method of cross attention is similar to the above attention calculation method, except that the query is obtained through linear mapping of image features, while the key and value are obtained through linear mapping of Object Query.

[0094] Based on the above embodiment, the detection results include the category of the corresponding predicted geometric primitive, the center point, the key point, the angle and distance of the key point relative to the center point, and the corresponding relationship between each sector area and the key point;

[0095] The primitive detection network is used to detect geometric primitives based on instance query features, and obtain the category, center point, and correspondence between each predicted geometric primitive and the corresponding sector area and key points;

[0096] The primitive detection network is further configured to determine, when a key point of any predicted geometric primitive exists in any sector-shaped area corresponding to the predicted geometric primitive, an angle and a distance of the key point of the predicted geometric primitive relative to the center point, and determine the key point of the predicted geometric primitive based on the angle and the distance of the key point of the predicted geometric primitive relative to the center point;

[0097] Among them, any prediction geometric primitive corresponds to multiple fan-shaped areas, and the central angles of the fan-shaped areas are the same.

[0098] Specifically, in an embodiment of the present invention, when the primitive detection network performs geometric primitive detection, it first uses instance query features to perform geometric primitive detection to obtain the category, center point and corresponding relationship between each predicted geometric primitive and the key point in the sample image.

[0099] Figure 3 This is an example diagram of the geometric primitive detection process provided by the present invention, such as Figure 3 As shown, after determining the center point of each predicted geometric primitive, its 360° range can be divided equally with the center point as the center of the circle. That is, the 360° range can be divided into N equal regions. Each region is a sector, so it can be called a sector region. Then, it is determined whether the key point of the predicted geometric primitive exists in each sector region. If so, the angle and distance of the key point of the predicted geometric primitive relative to the center point are further regressed to obtain the accurate key point information of the corresponding predicted geometric primitive, realizing the accurate acquisition of the geometric primitive parameters. Furthermore, the key point of the predicted geometric primitive can be calculated based on the angle and distance of the key point of the predicted geometric primitive relative to the center point obtained by regression.

[0100] For example, Figure 3 The predicted geometric primitive of the middle category is a triangle, which has three key points, namely the three vertices of the triangle, and corresponds to eight sector-shaped areas. The central angle of each sector-shaped area is 45°. It can be judged that there are key points in the 3rd, 5th, and 8th sector-shaped areas. The key points in this part of the sector-shaped area can be regressed to obtain the angle and distance of the key point relative to the center point, and the key point can be calculated based on this angle and distance.

[0101] In addition, it should be noted that after the model training is completed, geometric primitives with different numbers of key points can be adaptively detected, and geometric primitives with different numbers of key points in the sample image can be automatically located at the same time.

[0102] In the embodiment of the present invention, the geometric primitive detection method based on multiple key points in polar coordinates can adaptively detect various types of primitives without pre-defining the parameter form of the primitives, which is simpler, more efficient and reliable.

[0103] Based on the above embodiment, the sample label includes the category label, center point, key point, angle and distance of the key point relative to the center point of the corresponding sample geometric primitive, and the relationship label between each sector area and the key point; any sample geometric primitive corresponds to multiple sector areas, and the central angle of each sector area is the same; step 130 includes:

[0104] Determine the classification loss based on the category of each predicted geometric primitive, the correspondence between each sector area corresponding to each predicted geometric primitive and the key point, as well as the category label of each sample geometric primitive and the relationship label between each sector area corresponding to each sample geometric primitive and the key point;

[0105] Determine the regression loss based on the center point and key point of each predicted geometric primitive relative to the center point and the angle and distance of the center point and key point of each sample geometric primitive relative to the center point;

[0106] Determine the coincidence loss based on the center points and key points of each predicted geometric primitive and the center points and key points of each sample geometric primitive;

[0107] Based on classification loss, regression loss and coincidence loss, the parameters of the initial detection model are iterated to obtain the geometric primitive detection model.

[0108] Specifically, in step 130, the process of performing parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain the geometric primitive detection model may include:

[0109] Since the initial detection model in the embodiment of the present invention is built based on a single-stage end-to-end detection framework, the detection results obtained through model prediction include the category, center point, key point, angle and distance of the key point relative to the center point of each predicted geometric primitive in the sample image, as well as the corresponding relationship between each sector area and the key point; the sample label includes the category label, center point, key point, angle and distance of the key point relative to the center point of each sample geometric primitive in the sample image, as well as the relationship label between each sector area and the key point. Therefore, when training the network model, multiple loss functions are required to measure the gap between the predicted geometric primitive parameters and the sample labels. Specifically, for the category of the geometric primitive and whether there are key points in each sector area, Focal Loss can be used as the loss function to calculate the classification loss.

[0110] That is, the classification loss can be determined based on the category of each predicted geometric primitive, the corresponding relationship between each sector area and the key point corresponding to each predicted geometric primitive, the category label of each sample geometric primitive, and the relationship label between each sector area and the key point corresponding to each sample geometric primitive. Here, the classification loss FL(p t ) is calculated as follows:

[0111] FL(p t )=-α t (1-p t ) γ log(p t )

[0112] Where p t To predict the probability distribution of the category of geometric primitives, such as the probability that the geometric primitive belongs to a triangle, a quadrilateral, etc., or the probability that there are key points in the current fan-shaped area; α t is the weight coefficient of positive and negative samples, and γ is the adjustment factor.

[0113] Here, it should be noted that the prediction output of the initial detection model is usually not a certain category, but a category probability distribution, which represents the probability that the corresponding predicted geometric primitive belongs to each category (such as triangle, quadrilateral, etc.). The category of the predicted geometric primitive can be determined by the maximum probability in the category probability distribution. Similarly, the model output for whether there is a key point in the fan-shaped area is usually not a certain value, for example, corresponding or not corresponding, present or not, etc., but a probability, that is, the probability that the key point exists in the corresponding fan-shaped area. Based on this probability, the correspondence between the fan-shaped area and the key point can be determined. Positive and negative samples include positive samples and negative samples. Positive samples refer to sample geometric primitives whose detection results of the model are consistent with the sample labels, for example, sample geometric primitives whose categories detected by the model are consistent with the category labels in the sample labels, or sample geometric primitives whose correspondence between the fan-shaped area and the key point predicted by the model is consistent with the relationship labels in the sample labels.

[0114] For the regressed geometric primitive parameters, such as the angles and distances of the center point and key points relative to the center point, this embodiment of the present invention uses L1 loss as the loss function to calculate the regression loss. Specifically, the regression loss is determined based on the angles and distances of the center point and key points of the predicted geometric primitive relative to the center point, as well as the angles and distances of the center point and key points of each sample geometric primitive relative to the center point.

[0115] Here, the regression loss is calculated as follows:

[0116]

[0117] Where n represents the number of sample geometric primitives in the sample image, y i represents the sample label of the i-th sample geometric primitive in the sample image, It represents the detection result of the i-th predicted geometric primitive in the sample image.

[0118] For keypoint prediction, embodiments of the present invention can calculate the degree of overlap between the predicted geometric primitives and the sample geometric primitives based on the keypoints in the sample labels and the keypoints in the detection results. This can be used to determine the loss of the initial detection model in the geometric primitive detection task, i.e., the overlap loss. Specifically, the overlap loss can be determined based on the center point and keypoints of each predicted geometric primitive, as well as the center point and keypoints of each sample geometric primitive.

[0119] Furthermore, after obtaining the classification loss, regression loss, and overlap loss, the initial detection model can be trained based on these three factors to obtain a trained geometric primitive detection model. That is, the parameters of the initial detection model can be iterated based on the classification loss, regression loss, and overlap loss to adjust the parameters of the initial detection model. This allows the adjusted initial detection model to output detection results that are as consistent as possible with the sample labels when performing geometric primitive detection on sample images, thereby reducing its loss in the detection task and improving the model's performance in geometric primitive detection. By repeating the above process and through multiple rounds of parameter iteration, a geometric primitive detection model that performs extremely well in geometric primitive detection tasks can be obtained.

[0120] Based on the above embodiment, the overlap loss is determined based on the center point and key point of each predicted geometric primitive and the center point and key point of each sample geometric primitive, including:

[0121] Determining a degree of overlap between each predicted geometric primitive and the corresponding sample geometric primitive based on a distance between a center point of each sample geometric primitive and a key point, a distance between a center point of each sample geometric primitive and a key point of the corresponding predicted geometric primitive, a distance between a center point of each predicted geometric primitive and a key point, and a distance between a center point of each predicted geometric primitive and a key point of the corresponding sample geometric primitive;

[0122] Based on the degree of coincidence corresponding to each predicted geometric primitive, the coincidence loss is determined.

[0123] Specifically, the process of determining the overlap loss based on the center point and key point of each predicted geometric primitive and the center point and key point of each sample geometric primitive may include:

[0124] In order to ensure that the key points predicted by the model during model training are close to the real key points, in an embodiment of the present invention, the degree of overlap between each predicted geometric primitive output by the model and the corresponding sample geometric primitive, namely, PolygonIoU, can be calculated to measure the difference between the model in key point prediction and center point prediction in the geometric primitive detection task. Parameters can be adjusted accordingly to make the key points output by the adjusted model approach the real results.

[0125] Figure 4 This is an example diagram of the calculation of the degree of coincidence between the predicted geometric primitives and the sample geometric primitives provided by the present invention, such as Figure 4 As shown, the calculation of the overlap fully considers the distance between the center point of the geometric primitive and the key point, which can more accurately reflect the overlap between the predicted geometric primitive and the sample geometric primitive, thereby promoting the model to predict the key point more accurately.

[0126] The calculation formula of PolygonIoU is as follows:

[0127]

[0128] Where n represents the number of key points in the sample geometric primitives, Represents the distance between the center point of the sample geometric primitive and its i-th key point, Represents the distance between the center point of the sample geometric primitive and the i-th key point of the corresponding predicted geometric primitive, represents the distance between the center point of the predicted geometric primitive and its i-th key point, Represents the distance between the center point of the predicted geometric primitive and the i-th key point of its corresponding sample geometric primitive.

[0129] After obtaining the degree of overlap between each predicted geometric primitive and the corresponding sample geometric primitive, the overlap loss PolygonIoULoss of the initial detection model can be determined accordingly.

[0130] The calculation formula for coincidence loss is as follows:

[0131]

[0132] In the embodiment of the present invention, the polygon intersection-to-union ratio calculation method can more accurately reflect the degree of overlap between the predicted geometric primitives and the sample geometric primitives, and can further improve the accuracy of geometric primitive detection.

[0133] Based on the above embodiment, the sample image is determined based on the following steps:

[0134] Obtain an initial image set, where the initial image set includes multiple initial images;

[0135] Preprocessing is performed on each initial image in the initial image set to obtain a sample image set, where the sample image set includes multiple sample images; the preprocessing includes at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization.

[0136] Specifically, the process of determining the sample image may include:

[0137] It is understood that in the initial stage, a large dataset of geometric primitives is first constructed and then divided according to a set ratio to obtain a training dataset and a test dataset. To enrich the data diversity and multi-scale during model training, embodiments of the present invention can further process this dataset to obtain the final set of sample images used for model training. Therefore, the training dataset obtained by this division can be referred to as the initial image set. This initial image set contains multiple initial images.

[0138] Furthermore, each initial image in the initial image set can be preprocessed to obtain a preprocessed sample image set, and the sample image set contains multiple sample images; here, specifically, each initial image can be subjected to at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization to obtain each sample image.

[0139] Among them, multi-scale image scaling is to enlarge or reduce each initial image at multiple scales to provide multi-scale training data for the model, thereby increasing the model's robustness. Multi-dimensional image flipping is to randomly flip each initial image in multiple dimensions, such as up and down, left and right, and diagonally, to provide diverse training data for the model, thereby increasing the model's generalization ability. Pixel mean and variance normalization is to calculate the mean and variance of each pixel in the initial image set and use this mean and variance to normalize each initial image to accelerate the model's convergence and improve its stability.

[0140] Here, the formulas for pixel mean and variance normalization are as follows:

[0141]

[0142] Where x is the pixel of the initial image, μ is the mean, σ is the variance, and y is the normalized pixel.

[0143] Figure 5 This is the overall framework diagram of the geometric primitive detection model training method based on key points provided by the present invention, such as Figure 5 As shown, the method includes:

[0144] First, a sample image is determined, where the sample image contains multiple sample geometric primitives;

[0145] Then, based on the initial detection model, the multi-scale image features of the sample image are determined, the multi-scale image features are fused, and based on the fused instance query features, geometric primitive detection is performed to obtain the detection results of each predicted geometric primitive in the sample image;

[0146] Subsequently, based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, the parameters of the initial detection model are iterated to obtain a geometric primitive detection model.

[0147] Here, the sample image is determined based on the following steps: obtaining an initial image set, the initial image set containing multiple initial images; preprocessing each initial image in the initial image set to obtain a sample image set, the sample image set containing multiple sample images; the preprocessing includes at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization.

[0148] Among them, the initial detection model includes a feature extraction network, a feature fusion network, and a primitive detection network; the feature extraction network is used to perform multi-scale feature extraction on the sample image to obtain the multi-scale image features of the sample image; the feature fusion network is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features; the primitive detection network is used to perform geometric primitive detection based on instance query features to obtain the detection results of each predicted geometric primitive in the sample image.

[0149] Here, the feature fusion network includes an encoder and a decoder; the encoder is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features to obtain the fused features; the decoder is used to perform cross attention calculation based on the fused features and the initial query vector obtained by random initialization to obtain the instance query feature

[0150] Here, the detection results include the category, center point, key point, angle and distance of the key point relative to the center point of the corresponding predicted geometric primitive, as well as the corresponding relationship between each fan-shaped area and the key point; the primitive detection network is used to perform geometric primitive detection based on instance query features to obtain the category, center point and corresponding relationship between each fan-shaped area and the key point of each predicted geometric primitive; the primitive detection network is also used to determine the angle and distance of the key point of the predicted geometric primitive relative to the center point when the key point of the predicted geometric primitive exists in any fan-shaped area corresponding to any predicted geometric primitive, and determine the key point of the predicted geometric primitive based on the angle and distance of the key point of the predicted geometric primitive relative to the center point; wherein, any predicted geometric primitive corresponds to multiple fan-shaped areas, and the central angle of each fan-shaped area is the same.

[0151] Among them, the sample label includes the category label, center point, key point, angle and distance of the key point relative to the center point of the corresponding sample geometric primitive, and the relationship label between each sector area and the key point; any sample geometric primitive corresponds to multiple sector areas, and the central angle of each sector area is the same; based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, the initial detection model is iterated to obtain a geometric primitive detection model, including: based on the category of each predicted geometric primitive, the corresponding relationship between each sector area corresponding to each predicted geometric primitive and the key point, and The classification loss is determined based on the category label of each sample geometric primitive and the relationship label between each sector area and the key point corresponding to each sample geometric primitive; the regression loss is determined based on the center point of each predicted geometric primitive, the angle and distance of the key point relative to the center point, and the center point of each sample geometric primitive, the angle and distance of the key point relative to the center point; the coincidence loss is determined based on the center point and key point of each predicted geometric primitive, and the center point and key point of each sample geometric primitive; based on the classification loss, regression loss and coincidence loss, the parameters of the initial detection model are iterated to obtain the geometric primitive detection model.

[0152] Based on the center point and key point of each predicted geometric primitive, and the center point and key point of each sample geometric primitive, the overlap loss is determined, including: based on the distance between the center point of each sample geometric primitive and the key point, the distance between the center point of each sample geometric primitive and the key point of the corresponding predicted geometric primitive, the distance between the center point of each predicted geometric primitive and the key point, and the distance between the center point of each predicted geometric primitive and the key point of the corresponding sample geometric primitive, determining the overlap loss based on the corresponding overlap of each predicted geometric primitive.

[0153] In addition, it should be noted that after the entire training is completed, the testing process can be executed, calling the geometric primitive detection model to perform geometric primitive detection, thereby obtaining the detection results of each geometric primitive in each test image in the test dataset. The detection results here can be the category, key points, etc. of each geometric primitive in the test image.

[0154] In this embodiment of the present invention, the distribution of the training dataset is taken into account when determining the test dataset. This ensures that the training and test datasets have the same data distribution. This allows the model's training effectiveness and performance in geometric primitive detection to be verified in the post-training testing phase using the same test dataset. Specifically, a large geometric primitive dataset can be constructed and then divided into training and test datasets according to a predetermined ratio, such as 8:2 or 7:3, ensuring that the training and test datasets have the same data distribution.

[0155] The method provided by the embodiment of the present invention determines the multi-scale image features of the sample image through an initial detection model, fuses the multi-scale image features, and performs geometric primitive detection based on the instance query features obtained by the fusion to obtain the detection results of each predicted geometric primitive in the sample image; a geometric primitive detection model is trained based on the detection results of each predicted geometric primitive and the sample labels of each sample geometric primitive, thereby realizing automated geometric primitive detection based on the geometric primitive detection model, overcoming the defect of poor generalization performance of the geometric primitive detection method in the traditional scheme. The initial detection model is built through an end-to-end detection framework, and the geometric primitive detection model is obtained by training. This not only avoids the problem of difficulty in accurately representing geometric primitives and low detection accuracy in traditional target detection schemes, but also enhances the generalization performance of the network model and improves its stability and detection accuracy.

[0156] The following describes the key point-based geometric primitive detection model training device provided by the present invention. The key point-based geometric primitive detection model training device described below and the key point-based geometric primitive detection model training method described above can be referenced to each other.

[0157] Figure 6 Schematic diagram of the structure of the geometric primitive detection model training device based on key points provided by the present invention. Figure 6 As shown, the device includes:

[0158] A determining unit 610 is configured to determine a sample image, wherein the sample image includes a plurality of sample geometric primitives;

[0159] a detection unit 620 configured to determine multi-scale image features of the sample image based on an initial detection model, fuse the multi-scale image features, and perform geometric primitive detection based on the fused instance query features to obtain detection results for each predicted geometric primitive in the sample image;

[0160] The iteration unit 630 is configured to perform parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model.

[0161] The key point-based geometric primitive detection model training device provided by the present invention determines the multi-scale image features of the sample image through the initial detection model, fuses the multi-scale image features, and performs geometric primitive detection based on the instance query features obtained by the fusion to obtain the detection results of each predicted geometric primitive in the sample image; based on the detection results of each predicted geometric primitive and the sample labels of each sample geometric primitive, a geometric primitive detection model is trained to obtain automatic geometric primitive detection based on the geometric primitive detection model, overcoming the defect of poor generalization performance of the geometric primitive detection method in the traditional scheme. The initial detection model is built through the end-to-end detection framework, and the geometric primitive detection model is obtained by training. This not only avoids the problems of difficult accurate representation of geometric primitives and low detection accuracy in traditional target detection schemes, but also enhances the generalization performance of the network model and improves its stability and detection accuracy.

[0162] Based on the above embodiment, the initial detection model includes a feature extraction network, a feature fusion network, and a primitive detection network;

[0163] The feature extraction network is used to perform multi-scale feature extraction on the sample image to obtain multi-scale image features of the sample image;

[0164] The feature fusion network is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features;

[0165] The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain detection results of each predicted geometric primitive in the sample image.

[0166] Based on the above embodiment, the feature fusion network includes an encoder and a decoder;

[0167] The encoder is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features to obtain a fused feature;

[0168] The decoder is used to perform cross-attention calculation based on the fused features and an initial query vector obtained by random initialization to obtain instance query features.

[0169] Based on the above embodiment, the detection result includes the category of the corresponding predicted geometric primitive, the center point, the key point, the angle and distance of the key point relative to the center point, and the corresponding relationship between each sector area and the key point;

[0170] The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain the category and center point of each predicted geometric primitive and the corresponding relationship between each sector area and the key point;

[0171] The primitive detection network is further configured to determine, when a key point of any predicted geometric primitive exists in any sector-shaped area corresponding to any predicted geometric primitive, an angle and a distance of the key point of the predicted geometric primitive relative to the center point, and determine the key point of the predicted geometric primitive based on the angle and the distance of the key point of the predicted geometric primitive relative to the center point;

[0172] Among them, any prediction geometric primitive corresponds to multiple fan-shaped areas, and the central angles of the fan-shaped areas are the same.

[0173] Based on the above embodiment, the sample label includes the category label, center point, key point, angle and distance of the key point relative to the center point of the corresponding sample geometric primitive, and the relationship label between each corresponding sector area and the key point; any sample geometric primitive corresponds to multiple sector areas, and the central angle of each sector area is the same;

[0174] The iteration unit 630 is used to:

[0175] Determining a classification loss based on the category of each predicted geometric primitive, the correspondence between each sector area corresponding to each predicted geometric primitive and the key point, and the category label of each sample geometric primitive and the relationship label between each sector area corresponding to each sample geometric primitive and the key point;

[0176] Determining a regression loss based on the center point of each predicted geometric primitive, the angle and distance of the key point relative to the center point, and the angle and distance of the center point of each sample geometric primitive, the angle and distance of the key point relative to the center point;

[0177] Determining a coincidence loss based on the center points and key points of each predicted geometric primitive and the center points and key points of each sample geometric primitive;

[0178] Based on the classification loss, the regression loss and the coincidence loss, parameter iteration is performed on the initial detection model to obtain a geometric primitive detection model.

[0179] Based on the above embodiment, the iteration unit 630 is configured to:

[0180] Determining a degree of overlap between each predicted geometric primitive and the corresponding sample geometric primitive based on a distance between a center point of each sample geometric primitive and a key point, a distance between a center point of each sample geometric primitive and a key point of a corresponding predicted geometric primitive, a distance between a center point of each predicted geometric primitive and a key point, and a distance between a center point of each predicted geometric primitive and a key point of a corresponding sample geometric primitive;

[0181] Based on the overlap degrees corresponding to the predicted geometric primitives, an overlap loss is determined.

[0182] Based on the above embodiment, the sample image is determined based on the following steps:

[0183] Acquire an initial image set, wherein the initial image set includes a plurality of initial images;

[0184] Preprocessing is performed on each initial image in the initial image set to obtain a sample image set, where the sample image set contains multiple sample images; the preprocessing includes at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization.

[0185] Figure 7 An example of a physical structure diagram of an electronic device is shown below. Figure 7 As shown, the electronic device may include: a processor 710, a communication interface 720, a memory 730, and a communication bus 740, wherein the processor 710, the communication interface 720, and the memory 730 communicate with each other via the communication bus 740. The processor 710 may call the logic instructions in the memory 730 to execute a key point-based geometric primitive detection model training method, the method comprising: determining a sample image, wherein the sample image contains a plurality of sample geometric primitives; determining multi-scale image features of the sample image based on an initial detection model, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results of each predicted geometric primitive in the sample image; and performing parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample labels of each sample geometric primitive to obtain a geometric primitive detection model.

[0186] In addition, the logic instructions in the above-mentioned memory 730 can be implemented in the form of a software functional unit and can be stored in a computer-readable storage medium when sold or used as an independent product. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, or the part of the technical solution, can be embodied in the form of a software product. The computer software product is stored in a storage medium and includes several instructions for enabling a computer device (which can be a personal computer, server, or network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes: various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a magnetic disk or an optical disk.

[0187] On the other hand, the present invention also provides a computer program product, which includes a computer program stored on a non-transitory computer-readable storage medium, and the computer program includes program instructions. When the program instructions are executed by a computer, the computer can execute the key point-based geometric primitive detection model training method provided by the above methods, the method including: determining a sample image, wherein the sample image contains multiple sample geometric primitives; based on an initial detection model, determining multi-scale image features of the sample image, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results of each predicted geometric primitive in the sample image; based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, performing parameter iteration on the initial detection model to obtain a geometric primitive detection model.

[0188] On the other hand, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, is implemented to execute the key point-based geometric primitive detection model training method provided by the above-mentioned methods, the method comprising: determining a sample image, wherein the sample image contains multiple sample geometric primitives; based on an initial detection model, determining multi-scale image features of the sample image, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results of each predicted geometric primitive in the sample image; based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, performing parameter iteration on the initial detection model to obtain a geometric primitive detection model.

[0189] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, i.e., they may be located in one location or distributed across multiple network units. Some or all of the modules may be selected based on actual needs to achieve the objectives of the present embodiment. Persons of ordinary skill in the art will be able to understand and implement the present invention without inventive effort.

[0190] Through the description of the above embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus a necessary general hardware platform, or of course, by hardware. Based on this understanding, the essence of the above technical solution or the part that contributes to the existing technology can be embodied in the form of a software product. The computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a magnetic disk, an optical disk, etc., and includes a number of instructions for enabling a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in each embodiment or certain parts of the embodiments.

[0191] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit it. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present invention.

Claims

1. A method for training a geometric primitive detection model based on key points, characterized in that: include: Determine a sample image, wherein the sample image includes a plurality of sample geometric primitives; Determining multi-scale image features of the sample image based on the initial detection model, fusing the multi-scale image features, and performing geometric primitive detection based on the fused instance query features to obtain detection results of each predicted geometric primitive in the sample image; Based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive, the initial detection model is iterated on parameters to obtain a geometric primitive detection model.

2. The method for training a geometric primitive detection model based on key points according to claim 1, wherein: The initial detection model includes a feature extraction network, a feature fusion network, and a primitive detection network; The feature extraction network is used to perform multi-scale feature extraction on the sample image to obtain multi-scale image features of the sample image; The feature fusion network is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features; The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain detection results of each predicted geometric primitive in the sample image.

3. The method for training a geometric primitive detection model based on key points according to claim 2, wherein: The feature fusion network includes an encoder and a decoder; The encoder is used to fuse the multi-scale image features based on the correlation between each image feature in the multi-scale image features to obtain a fused feature; The decoder is used to perform cross-attention calculation based on the fused features and an initial query vector obtained by random initialization to obtain instance query features.

4. The method for training a geometric primitive detection model based on key points according to claim 2 or 3, wherein: The detection results include the category, center point, key point, angle and distance of the key point relative to the center point of the corresponding predicted geometric primitive, and the corresponding relationship between each sector area and the key point; The primitive detection network is used to perform geometric primitive detection based on the instance query feature to obtain the category and center point of each predicted geometric primitive and the corresponding relationship between each sector area and the key point; The primitive detection network is further configured to determine, when a key point of any predicted geometric primitive exists in any sector-shaped area corresponding to any predicted geometric primitive, an angle and a distance of the key point of the predicted geometric primitive relative to the center point, and determine the key point of the predicted geometric primitive based on the angle and the distance of the key point of the predicted geometric primitive relative to the center point; Among them, any prediction geometric primitive corresponds to multiple fan-shaped areas, and the central angles of the fan-shaped areas are the same.

5. The method for training a geometric primitive detection model based on key points according to claim 4, wherein: The sample labels include the category label, center point, key point, angle and distance of the key point relative to the center point of the corresponding sample geometric primitive, and the relationship label between each sector area and the key point; any sample geometric primitive corresponds to multiple sector areas, and the central angle of each sector area is the same; The method of performing parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model includes: Determining a classification loss based on the category of each predicted geometric primitive, the correspondence between each sector area corresponding to each predicted geometric primitive and the key point, and the category label of each sample geometric primitive and the relationship label between each sector area corresponding to each sample geometric primitive and the key point; Determining a regression loss based on the center point of each predicted geometric primitive, the angle and distance of the key point relative to the center point, and the angle and distance of the center point of each sample geometric primitive, the angle and distance of the key point relative to the center point; Determining a coincidence loss based on the center points and key points of each predicted geometric primitive and the center points and key points of each sample geometric primitive; Based on the classification loss, the regression loss and the coincidence loss, parameter iteration is performed on the initial detection model to obtain a geometric primitive detection model.

6. The method for training a geometric primitive detection model based on key points according to claim 5, wherein: The determining of the overlap loss based on the center point and the key point of each predicted geometric primitive and the center point and the key point of each sample geometric primitive includes: Determining a degree of overlap between each predicted geometric primitive and the corresponding sample geometric primitive based on a distance between a center point of each sample geometric primitive and a key point, a distance between a center point of each sample geometric primitive and a key point of a corresponding predicted geometric primitive, a distance between a center point of each predicted geometric primitive and a key point, and a distance between a center point of each predicted geometric primitive and a key point of a corresponding sample geometric primitive; Based on the overlap degrees corresponding to the predicted geometric primitives, an overlap loss is determined.

7. The method for training a geometric primitive detection model based on key points according to any one of claims 1 to 3, characterized in that: The sample image is determined based on the following steps: Acquire an initial image set, wherein the initial image set includes a plurality of initial images; Preprocessing is performed on each initial image in the initial image set to obtain a sample image set, where the sample image set contains multiple sample images; the preprocessing includes at least one of multi-scale image scaling, multi-dimensional image flipping, pixel mean and variance normalization.

8. A geometric primitive detection model training device based on key points, characterized in that: include: a determining unit, configured to determine a sample image, wherein the sample image includes a plurality of sample geometric primitives; a detection unit, configured to determine multi-scale image features of the sample image based on an initial detection model, fuse the multi-scale image features, and perform geometric primitive detection based on the fused instance query features to obtain a detection result of each predicted geometric primitive in the sample image; The iterative unit is used to perform parameter iteration on the initial detection model based on the detection results of each predicted geometric primitive and the sample label of each sample geometric primitive to obtain a geometric primitive detection model.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the program, the key point-based geometric primitive detection model training method according to any one of claims 1 to 7 is implemented.

10. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the key point-based geometric primitive detection model training method according to any one of claims 1 to 7 is implemented.

Images