Classification model training method and device, image classification method and device, equipment and medium
By using branch regression network in the object detection model to supervise and train the backbone classification network, and abandoning the branch regression network after training, the problem of inconsistent confidence calculation in the object detection model is solved, and the accuracy of classification results is improved.
Patent Information
- Application Number
- CN202411678008.3
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-21
- Publication Date
- 2025-03-28
AI Technical Summary
The confidence calculation of the object detection model is not uniform during the training and application process, resulting in a decrease in the accuracy of the classification results.
The training images are targeted through the backbone classification network of the basic model, and feature images are obtained from the backbone classification network through the branch regression network for object detection, realizing the supervision and training of the backbone classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned and the backbone classification network is used as the target classification model.
This avoids the problem of inconsistent confidence calculations during training and application, and effectively improves the accuracy of classification results.
Smart Images

Figure CN119723154A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of deep learning technology, and in particular to a classification model training method, an image classification method, an apparatus, a computer device and a storage medium. Background Art
[0002] Target classification and target detection are two common tasks in computer vision. Target classification tasks do not care about the location of the target. However, some classification tasks hope to consider the location information of the detection target while performing target classification. For example, the color recognition of pedestrian shoulder bags in the detection image, the recognition of pedestrians holding mobile phones, and the recognition of road diseases. Taking the classification task of pedestrian shoulder bag color recognition as an example, when classifying shoulder bags, it is necessary to first determine whether the shoulder bag is on the person, and then output the color category corresponding to the shoulder bag. When the pedestrian does not have a shoulder bag, or the shoulder bag is placed somewhere else (such as on a clothes rack) instead of on the person, the detection image will be output as a background image.
[0003] At present, target classification based on location information is usually based on target detection models, such as YOLO, DETR, etc., to achieve target classification based on the detected target location information. When the target detection model outputs, the position output is discarded and only the classification result is retained.
[0004] However, in the process of training the target detection model, it is usually necessary to fuse the probability of whether the detection image is a background class, the probability of the detection image belonging to a certain color, and the IoU (intersection over union) of the predicted box and the true box of the detection image to calculate the confidence. In the process of applying the target detection model, since there is no true box, the confidence can only be calculated by fusing the probability of whether the detection image is a background class and the probability of the detection image belonging to a certain color, so as to eliminate the classification results with low confidence. It can be seen that there is an inconsistency in the confidence calculation during the target detection model training process and the model application process, which will directly affect the output results of the model and cause the classification results output by the target detection model to be less accurate. Summary of the invention
[0005] Embodiments of the present invention provide a classification model training method, an image classification method, an apparatus, a computer device, and a storage medium to solve the problem of reduced accuracy of classification results output by a target detection model.
[0006] A classification model training method, the method comprising: Acquire training images, each of which is annotated with a category label and a position label of a detection target; Inputting the training image into a basic model to iteratively train the basic model, wherein the basic model includes a trunk classification network and a branch regression network, wherein the trunk classification network is used to perform target classification on the training image, and the branch regression network is used to obtain a feature image from the trunk classification network to perform target detection on the training image; After the basic model is trained, the branch regression network is discarded and the trunk classification network is determined as the target classification model.
[0007] In the above method, optionally, the step of inputting the training image into a basic model to iteratively train a trunk classification network and a branch regression network of the basic model comprises: Inputting the training image into the basic model to obtain the predicted classification result output by the trunk classification network and the first predicted coordinate output by the branch regression network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, the total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
[0008] In the above method, optionally, the branch regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
[0009] A method for image classification, comprising: Get the target image; Inputting the target image into the target classification model to obtain a probability value that the target image belongs to each classification category, wherein the target classification model is trained by the classification model training method according to any one of claims 1 to 3; Determine that the classification category with the maximum probability value is the target category of the target image.
[0010] In the above method, optionally, acquiring the target image includes: Get the original image; Inputting the original image into a target detection model to obtain at least one second predicted coordinate and a prediction confidence of the second predicted coordinate; Determining whether the prediction confidence exceeds a preset confidence threshold; If the prediction confidence exceeds a preset confidence threshold, the target image is captured from the second prediction coordinates.
[0011] A classification model training device, comprising: A training image acquisition module, used to acquire training images, each of which is annotated with a category label and a position label of a detection target; A classification model training module, used for inputting the training image into a basic model to iteratively train the basic model, wherein the basic model includes a trunk classification network and a branch regression network, and the branch regression network obtains a feature image from the trunk classification network; The classification model determination module is used to discard the branch regression network and determine the trunk classification network as the target classification model after the basic model is trained.
[0012] In the above device, optionally, the classification model training module is used to: Inputting the training image into a basic model, obtaining a predicted classification result output by a trunk classification network of the basic model and a first predicted coordinate output by a branch regression network, wherein the branch regression network obtains a feature image from the trunk classification network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, the total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
[0013] In the above device, optionally, the branch regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
[0014] A computer device comprises a memory, a processor and a computer program stored in the memory and executable on the processor, wherein when the processor executes the computer program, the classification model training method as described in any one of the above items is implemented.
[0015] A computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the classification model training method as described in any one of the above items is implemented.
[0016] The above-mentioned classification model training method, image classification method, device, computer equipment and storage medium perform target classification on the training image through the trunk classification network of the basic model, obtain feature images from the trunk classification network through the branch regression network of the basic model, and perform target detection on the training image, so as to realize the supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as the target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing classification tasks based on position information through the target detection model, it avoids the problem of inconsistent confidence calculation in the model training process and the execution process, and effectively improves the accuracy of the classification result. BRIEF DESCRIPTION OF THE DRAWINGS
[0017] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments of the present invention will be briefly introduced below. Obviously, the accompanying drawings in the following description are only some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying creative labor.
[0018] Figure 1 This is a flow chart for implementing a classification model training method in one embodiment of the present invention; Figure 2 is a schematic diagram of a network structure of a basic model in one embodiment of the present invention; Figure 3 It is a partial implementation flow chart of the classification model training method in one embodiment of the present invention; Figure 4 It is a partial implementation flow chart of the classification model training method in one embodiment of the present invention; Figure 5 is a flow chart of an implementation of an image classification method in one embodiment of the present invention; Figure 6 is a partial flow chart of an image classification method in one embodiment of the present invention; Figure 7is a structural diagram of a classification model training device in one embodiment of the present invention; Figure 8 It is a structural schematic diagram of a computer device in one embodiment of the present invention. DETAILED DESCRIPTION
[0019] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.
[0020] It should be understood that when used in the present specification and the appended claims, the term "comprising" indicates the presence of described features, integers, steps, operations, elements and / or components, but does not exclude the presence or addition of one or more other features, integers, steps, operations, elements, components and / or combinations thereof.
[0021] It should also be understood that the term "and / or" used in the present description and the appended claims refers to and includes any and all possible combinations of one or more of the associated listed items.
[0022] As used in the present specification and the appended claims, the term "if" may be interpreted as "when" or "upon" or "in response to determining" or "in response to detecting", depending on the context. Similarly, the phrase "if it is determined" or "if [described condition or event] is detected" may be interpreted as meaning "upon determination" or "in response to determining" or "upon detection of [described condition or event]" or "in response to detecting [described condition or event]", depending on the context.
[0023] In addition, in the description of the present specification and the appended claims, the terms "first", "second", "third", etc. are only used to distinguish the descriptions and cannot be understood as indicating or implying relative importance.
[0024] References to "one embodiment" or "some embodiments" etc. described in the present specification mean that one or more embodiments of the present invention include specific features, structures or characteristics described in conjunction with the embodiment. Therefore, the phrases "in one embodiment", "in some embodiments", "in some other embodiments", "in some other embodiments", etc. appearing in different places in this specification do not necessarily refer to the same embodiment, but mean "one or more but not all embodiments", unless otherwise specifically emphasized in other ways. The terms "including", "comprising", "having" and their variations all mean "including but not limited to", unless otherwise specifically emphasized in other ways.
[0025] The present invention discloses a classification model training method, an image classification method, an apparatus, a computer device and a storage medium. The training image is subjected to target classification through the trunk classification network of the basic model, and the feature image is obtained from the trunk classification network through the branch regression network of the basic model to perform target detection on the training image, thereby realizing the supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as the target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing the classification task through the target detection model, it avoids the problem of inconsistent confidence calculation in the training process and the execution process, and effectively improves the accuracy of the classification result.
[0026] In one embodiment, Figure 1 FIG. 1 is a flowchart of a classification model training method disclosed in this embodiment. The method is applicable to electronic devices with image processing capabilities, such as mobile phones, tablet computers, laptops, personal computers, servers, and other electronic devices. The method specifically includes the following steps: S101: Acquire training images, each of which is annotated with a category label and a position label of a detection target.
[0027] The training data in this embodiment includes a large amount of image data, and each image data is annotated with a category label and a position label.
[0028] In a specific implementation, the training images in this embodiment may include images containing the detection target and the reference target of the detection target, images containing only the detection target, images containing only the reference target, and images not containing the detection target and the reference target. Among them, the reference target is used to characterize the position information of the detection target. When the training image contains both the detection target and the detection target, the category label annotated on the training image is the specific category of the detection target, and when the training image contains only the image of the detection target, only the image of the reference target, and does not contain the detection target and the reference target, the category label annotated on the training image is the background class.
[0029] In addition, the position label of the training image in this embodiment is used to characterize the position information of the detection target of the training image, such as the coordinate data of the upper left corner and the coordinate data of the lower right corner of the detection box of the detection target, or the coordinate data of the lower left corner and the coordinate data of the upper right corner of the detection box of the detection target.
[0030] For example, take the classification task of the bag category on a pedestrian as an example, where the bag is the detection target, the pedestrian is the reference target, the location information of the bag is on the pedestrian, and the bag category includes but is not limited to one or more types of shoulder bags, backpacks, messenger bags, computer bags, and wallets. When the training image contains pedestrians and the bag is on the pedestrian, the category label of the training image is the specific category of the bag, and the location label of the training image is the specific coordinates of the detection target. When the training image only contains pedestrians, or only contains bags, or does not contain pedestrians and babies, the category label of the training image is the background class, and the location label of the training image is a unified invalid coordinate, such as (0,0).
[0031] S102: Input the training image into the basic model to iteratively train the basic model.
[0032] Among them, the basic model includes a trunk classification network and a branch regression network. The trunk classification network is used to classify targets on training images, and the branch regression network is used to obtain feature images from the trunk classification network to perform target detection on training images.
[0033] like Figure 2As shown, it is a schematic diagram of the structure of the basic model in this embodiment, wherein the input end of the backbone classification network is used to input the training image, Stage1, Stage2, Stage3 and Stage4 are feature extraction networks of different feature extraction stages, Stage1, Stage2, Stage3 and Stage4 are used to extract features from the training image in sequence to obtain feature images, and then output the classification results through the global pooling layer and the fully connected layer, wherein the shapes of the feature images output by Stage1, Stage2, Stage3 and Stage4 of the backbone classification network can be different in terms of the number of channels, height and width. Figure 2 As shown in , the height and width of the feature image output by Stage 1 are , the height and width of the feature image output by Stage2 are , the height and width of the feature image output by Stage3 are , the height and width of the feature image output by Stage4 are . The classification branch network obtains feature images from Stage 1, Stage 2, and Stage 3 of the trunk classification network respectively, and then performs global pooling and splicing through the global pooling layer based on the three feature images obtained, and finally outputs the regression result through the fully connected layer. It should be noted that the structural diagram of the above-mentioned basic model is only an exemplary description of the basic model in this embodiment, and it does not limit the regression branch in this embodiment to only obtain feature images from Stage 1, Stage 2, and Stage 3 of the trunk classification network. The regression branch in this embodiment can obtain feature images from any one or more feature extraction stages in Stage 1, Stage 2, Stage 3, and Stage 4 of the trunk classification network. Therefore, the specific structure of the basic model is not limited in this embodiment.
[0034] Furthermore, step S102 in this embodiment can be implemented by the following steps: Figure 3 As shown: S301: Input the training image into the basic model to obtain the predicted classification result output by the trunk classification network and the first predicted coordinate output by the branch regression network.
[0035] The training image is input into the basic model, the feature image is extracted through the backbone classification network, the feature image is classified, the feature image obtained from the backbone classification network is obtained, the target is detected on the feature image through the branch regression network, and the first predicted coordinate is obtained. In other words, during the basic model training process, the basic model will have two outputs, namely the predicted classification result and the first predicted coordinate.
[0036] S302: Calculate a classification loss value based on the predicted classification result and the category label, and calculate a regression loss value based on the first predicted coordinate and the position label.
[0037] In a specific implementation, the classification loss value in this embodiment includes but is not limited to being calculated based on any one of the loss functions such as the cross entropy loss function (CE-Loss), the focal loss function (Focal-Loss), etc. The specific calculation method of the classification loss value is not limited in this embodiment.
[0038] In one implementation, when the classification loss value is calculated based on the cross entropy loss function, the calculation formula of the cross entropy loss function can be as follows: in, Indicates the first The value of the class, which is 0 or 1. Indicates the first The smaller the value of the cross entropy loss function, the smaller the difference between the model prediction and the true label, that is, the better the performance of the model.
[0039] In addition, the regression loss value in this embodiment includes but is not limited to being calculated based on any one of the loss functions including the mean square error loss function (MSE-Loss), the absolute error loss function (MAE-Loss) and the intersection over union loss function (IoU-Loss).
[0040] In one implementation, when the regression loss value is calculated based on the cross entropy loss function, the calculation formula of the cross entropy loss function can be as follows: in, No. The class labels of the training images, Indicates The predicted classification results of training images, represents the number of training images, Represents the regression loss value.
[0041] It should be noted that the above is only a partial implementation method of calculating the classification loss value and the regression loss value in this embodiment. In this embodiment, the classification loss value and the regression loss value can also be calculated by other loss functions. In this embodiment, the method for obtaining the classification loss value and the regression loss value is not specifically limited.
[0042] The predicted classification result and category label are input into any loss function such as the cross entropy loss function and the focal loss function to obtain the classification loss value; after the first predicted coordinate and position label are normalized, they are input into any loss function such as the mean square error loss function, the absolute error loss function and the intersection-over-union loss function to calculate the regression loss value.
[0043] S303: Based on the classification loss value and the regression loss value, a total loss value is calculated.
[0044] The classification loss value and regression loss value are weighted summed to get the total loss value.
[0045] In a specific implementation, the classification loss value and the regression loss value are input into a weighted summation formula to obtain a total loss value, where the weighted summation formula can be shown as follows: in, represents the classification loss value, represents the regression loss value, Represents hyperparameters, which can be set according to actual needs. Represents the total loss value.
[0046] S304: When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
[0047] If the total loss value is greater than the loss threshold, it is determined that the basic model is fully trained; if the total loss value is less than or equal to the loss threshold, the basic model continues to be iteratively trained based on the training image, and the model parameters of the basic model are continuously updated during the training process until the total loss value is greater than the loss threshold, and it is determined that the basic model is fully trained.
[0048] S103: After the basic model is trained, the branch regression network is discarded, and the trunk classification network is determined as the target classification model.
[0049] That is to say, the branch regression network is only used in the training process of the basic model to supervise the training of the trunk classification network. After the basic model is trained, the branch regression network no longer needs to perform target detection, and only needs to output the predicted classification results through the trunk classification network. Therefore, the branch regression network is discarded and the trunk classification network is determined as the target classification model.
[0050] In summary, a classification model training method disclosed in this embodiment performs target classification on training images through the trunk classification network of the basic model, obtains feature images from the trunk classification network through the branch regression network of the basic model, and performs target detection on the training images, thereby realizing the supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as the target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing the classification task through the target detection model, it avoids the problem of inconsistent confidence calculation in the training process and the execution process, and effectively improves the accuracy of the classification result.
[0051] In one embodiment, the branched regression network in this embodiment outputs the predicted position result in the following manner: Figure 4 As shown: S401: Obtain a feature image extracted by at least one feature extraction stage of a backbone classification network.
[0052] It can be understood that there is no feature extraction network in the branch regression network in this embodiment. Therefore, it is necessary to share the feature image extracted by the trunk classification network. Therefore, it is necessary to obtain the feature image extracted from at least one feature extraction stage of the trunk classification network, and then perform the subsequent steps of target detection based on the acquired feature image.
[0053] S402: unify the number of channels of each feature image, and perform global pooling to obtain a standard feature image corresponding to each feature image.
[0054] It should be understood that the backbone classification network includes multiple feature extraction stages, and the number of channels, height and width of the feature images extracted in each feature extraction stage may be different. Therefore, after obtaining multiple feature images from the feature extraction stages of multiple backbone classification networks, it is necessary to first unify the number of channels of each feature image, and then perform global pooling on the feature images with the same number of channels to obtain standard images of the same shape.
[0055] In a specific implementation, in this embodiment, the number of channels of each feature image can be unified through 1×1 convolution, and then the feature images with the same number of channels are input into the global pooling layer to obtain a standard image with the same shape.
[0056] For example, the shapes obtained in the three feature extraction stages are , , As an example, after 1x1 convolution, the shapes obtained are , , After global pooling, the feature images are , , The standard image.
[0057] S403: splicing the standard feature images to obtain the target feature image.
[0058] In a specific implementation, in this embodiment, the standard images can be spliced in the channel dimension to obtain the target image.
[0059] For example, the shapes are , , Taking the standard image of as an example, after splicing the standard image in the channel dimension, the shape of the target image is .
[0060] S404: Input the target feature image into the position prediction layer to obtain the first predicted coordinates of the training image.
[0061] Among them, the position prediction layer is a fully connected layer, which is used to predict the position of the detection target in the target feature image to obtain the first predicted coordinates of the training image.
[0062] In a specific implementation, in this embodiment, the target feature image can be flattened to obtain a feature vector of the target feature image, and then the feature vector is input into the position prediction layer to obtain the first predicted coordinates of the training image. The first predicted coordinates include at least two coordinates, based on which a prediction box can be determined, and the two coordinates can be the upper left corner coordinates and the lower right corner coordinates of the prediction box, or the two coordinates can be the lower left corner coordinates and the upper right corner coordinates of the prediction box.
[0063] like Figure 5 FIG. 1 is a flowchart of an implementation of an image classification method disclosed in an embodiment of the present invention, and the method includes the following steps: S501: Acquire target image:.
[0064] S502: Input the target image into the target classification model to obtain the probability value of the target image belonging to each classification category.
[0065] The target classification model is obtained by training using the classification model training method disclosed in any of the above embodiments.
[0066] S503: Determine the classification category with the maximum probability value as the target category of the target image.
[0067] The target image is an image that needs to be classified by the target classification model in practical applications after the target classification model is trained.
[0068] The target image is input into the target classification model. First, the target image is feature extracted to obtain a feature image. Then, the feature image is flattened to obtain a feature vector. The feature vector is input into the fully connected layer to obtain the probability value of the target image belonging to each classification category. The larger the probability value, the higher the probability that the target image belongs to that classification category. The classification category with the largest probability value is selected as the target category of the target image.
[0069] For example, taking the classification task of the bag category on pedestrians as an example, the bag categories include single-shoulder bags, double-shoulder bags, and messenger bags. The image containing the pedestrian and the bag is input into the target classification model, and the probability of the bag in the image being a single-shoulder bag is 0.05, the probability of a double-shoulder bag is 0.85, and the probability of a messenger bag is 0.10. It can be seen that the probability of a double-shoulder bag is the highest. Therefore, the double-shoulder bag is finally determined as the final output result.
[0070] To sum up, the target classification model in this embodiment calculates the confidence in a unified manner during the training process and the actual application process, that is, the probability value of each classification category, which avoids the problem of inconsistent confidence calculation of the target detection model during the training process and the actual application process. Therefore, the target classification model in this embodiment can have better classification accuracy.
[0071] based on Figure 5 In the specific implementation, step S501 in this embodiment is implemented by the following steps: Figure 6 As shown: S601: Acquire an original image.
[0072] In a specific implementation, the original image in this embodiment may be an image acquired through big data, or an image captured by a device with a shooting function, such as a mobile phone, a camera, a camcorder, and other electronic devices such as a camera.
[0073] S602: Input the original image into the target detection model to obtain at least one second predicted coordinate and the prediction confidence of the second predicted coordinate.
[0074] It should be understood that the acquired original image may include more than one detection target, and therefore, it is necessary to separate each detection target from the original image into a target image, and then classify the target images.
[0075] In a specific implementation, the target detection model in this embodiment includes but is not limited to any one of the YOLO series models, DETR series models, etc. After the original image is input into the target detection model, the target detection model outputs at least one second predicted coordinate and the prediction confidence of the second predicted coordinate.
[0076] S603: Determine whether the prediction confidence exceeds a preset confidence threshold.
[0077] It should be noted that since the target detection model can only calculate the confidence based on the probability of whether the image in the second predicted coordinate is the background type and the probability of belonging to the classification category, it may lead to low output confidence accuracy. Therefore, the confidence threshold needs to be set to a lower value, such as 0.6, 0.65, etc., to obtain as many target images as possible.
[0078] If the prediction confidence exceeds the confidence threshold, the target image is captured from the second prediction coordinates. If the prediction confidence is less than or equal to the confidence threshold, the second prediction coordinates are discarded, that is, the target image is no longer captured from the second prediction coordinates.
[0079] S604: If the prediction confidence exceeds a preset confidence threshold, a target image is captured from the second prediction coordinates.
[0080] In a specific implementation, in this embodiment, the original image can be cropped by an image processing library (such as OpenCV, PIL, etc.). The cropping operation will cut out the target image from the original image based on the coordinates and size of the second predicted coordinates.
[0081] It can be understood that in this embodiment, each target image that may be a detection target can be cut out from the original image through the target detection model, and then the target image can be input into the target classification model to classify the detection target in the target image. For example, taking the classification task as the bag category on the pedestrian as an example, each target image containing a bag is detected from the original image through the target detection model, and then the target image is input into the target classification model to classify the target image through the target classification model. If there is only a bag in the target image, the probability value of the target image belonging to the background class is the largest. If there is no bag in the target image, the probability value of the target image belonging to the background class is the largest. If there is a bag in the target image and the bag is on the pedestrian, the color category with the largest probability value of the bag in the target image is output. Thus, through the target detection model and the target classification model, target classification of complex images can be achieved.
[0082] To sum up, in this embodiment, the target detection model can not only extract the target image from the original image, but also perform a preliminary screening of the target image, that is, not extracting the target image from the second prediction coordinate whose prediction confidence is less than or equal to the confidence threshold, which can reduce the number of invalid images in the target image and is conducive to improving the classification efficiency of the target classification model.
[0083] It should be understood that the order of execution of the steps in the above embodiment does not necessarily mean the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiment of the present invention.
[0084] In one embodiment, Figure 7 FIG. 1 is a schematic diagram of the structure of a classification model training device disclosed in this embodiment. The device includes a training image acquisition module, a classification model training module and a classification model determination module. Each model in this embodiment is used for: The training image acquisition module 701 is used to acquire training images, each of which is annotated with a category label and a position label of a detection target; The classification model training module 702 is used to input the training image into the basic model to iteratively train the basic model, wherein the basic model includes a trunk classification network and a branch regression network, and the branch regression network obtains feature images from the trunk classification network; The classification model determination module 703 is used to discard the branch regression network after the basic model is trained and determine the trunk classification network as the target classification model.
[0085] The present invention discloses a classification model training device, which performs target classification on training images through a trunk classification network of a basic model, obtains feature images from the trunk classification network through a branch regression network of the basic model, and performs target detection on the training images, thereby realizing supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as the target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing the classification task through the target detection model, it avoids the problem of inconsistent confidence calculation in the training process and the execution process, and effectively improves the accuracy of the classification result.
[0086] In one implementation, the classification model training module 701 is used to: The training image is input into the basic model to obtain the predicted classification result output by the trunk classification network of the basic model and the first predicted coordinate output by the branch regression network, and the branch regression network obtains the feature image from the trunk classification network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinate and the position label; Based on the classification loss value and regression loss value, the total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the base model is fully trained.
[0087] In one implementation, the branched regression network outputs the predicted position result in the following manner: Obtaining a feature image extracted by at least one feature extraction stage of a backbone classification network; Unify the number of channels of each feature image and perform global pooling to obtain the standard feature image corresponding to each feature image; The standard feature images are spliced to obtain the target feature image; The target feature image is input into the position prediction layer to obtain the first predicted coordinates of the training image.
[0088] In one implementation, the target classification model classifies the target image as follows: Input the target image into the target classification model to obtain the probability value of the target image belonging to each classification category; Determine the classification category with the maximum probability value as the target category of the target image.
[0089] In one implementation, the target image is obtained by: Get the original image; Inputting the original image into the target detection model to obtain at least one second predicted coordinate and a prediction confidence of the second predicted coordinate; Determine whether the prediction confidence exceeds a preset confidence threshold; If the prediction confidence exceeds a preset confidence threshold, the target image is captured from the second prediction coordinates.
[0090] For the specific limitations of the classification model training device, please refer to the relevant limitations of the classification model training method above, which will not be repeated here. Each module in the above-mentioned classification model training device can be implemented in whole or in part by software, hardware and a combination thereof. The above-mentioned modules can be embedded in or independent of the processor in the computer device in the form of hardware, or can be stored in the memory of the computer device in the form of software, so that the processor can call and execute the operations corresponding to the above modules.
[0091] In one embodiment, the present application discloses a computer device, which may be a server, and its internal structure diagram may be as follows: Figure 8As shown. The computer device includes a processor, a memory, a network interface and a database connected through a system bus. Among them, the processor of the computer device is used to provide computing and control capabilities. The memory of the computer device includes a non-volatile storage medium and an internal memory. The non-volatile storage medium stores an operating system, a computer program and a database. The internal memory provides an environment for the operation of the operating system and the computer program in the non-volatile storage medium. The network interface of the computer device is used to communicate with an external terminal through a network connection. When the computer program is executed by the processor, a classification model training method is implemented.
[0092] In one embodiment, a computer device is provided, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the following steps when executing the computer program: Obtain training images, each of which is annotated with a category label and a position label of the detection target; Inputting the training image into the basic model to iteratively train the basic model, the basic model includes a trunk classification network and a branch regression network, the trunk classification network is used to classify the target of the training image, and the branch regression network is used to obtain the feature image from the trunk classification network to perform target detection on the training image; After the basic model is trained, the branch regression network is discarded and the trunk classification network is determined as the target classification model.
[0093] The present invention discloses a computer device, which performs target classification on training images through a trunk classification network of a basic model, obtains feature images from the trunk classification network through a branch regression network of the basic model, and performs target detection on the training images, thereby realizing supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as the target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing the classification task through the target detection model, it avoids the problem of inconsistent confidence calculation in the training process and the execution process, and effectively improves the accuracy of the classification result.
[0094] In one implementation, inputting the training image into a basic model to iteratively train a trunk classification network and a branch regression network of the basic model includes: Inputting the training image into the basic model to obtain the predicted classification result output by the trunk classification network and the first predicted coordinate output by the branch regression network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, the total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
[0095] In one implementation, the branched regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
[0096] In one embodiment, this embodiment discloses a computer-readable storage medium, when the instructions in the computer-readable storage medium are executed by a processor in a computer device, the computer device can perform the steps of any embodiment of a classification model training method disclosed in the present invention. The computer-readable storage medium can be non-volatile or volatile.
[0097] In one embodiment, a computer readable storage medium is provided, on which a computer program is stored, and when the computer program is executed by a processor, the following steps are implemented: Obtain training images, each of which is annotated with a category label and a position label of the detection target; Inputting the training image into the basic model to iteratively train the basic model, the basic model includes a trunk classification network and a branch regression network, the trunk classification network is used to classify the target of the training image, and the branch regression network is used to obtain the feature image from the trunk classification network to perform target detection on the training image; After the basic model is trained, the branch regression network is discarded and the trunk classification network is determined as the target classification model.
[0098] The present invention discloses a computer-readable storage medium. After being executed by a processor in a computer device, the computer-readable storage medium performs target classification on a training image through a trunk classification network of a basic model, obtains a feature image from the trunk classification network through a branch regression network of the basic model, and performs target detection on the training image, thereby realizing supervised training of the trunk classification network by the branch regression network. After the basic model training is completed, the branch regression network is abandoned, and the trunk classification network is used as a target classification model. It can be seen that the present invention fully considers the position information of the detection target in the training image through the branch regression network, and by abandoning the branch regression network and using the trunk classification network as the target classification model, compared with performing the classification task through the target detection model, it avoids the problem of inconsistent confidence calculation in the training process and the execution process, and effectively improves the accuracy of the classification result.
[0099] In one implementation, inputting the training image into a basic model to iteratively train a trunk classification network and a branch regression network of the basic model includes: Inputting the training image into the basic model to obtain the predicted classification result output by the trunk classification network and the first predicted coordinate output by the branch regression network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, the total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
[0100] In one implementation, the branched regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
[0101] Those of ordinary skill in the art can understand that all or part of the processes in the above-mentioned embodiment methods can be completed by instructing the relevant hardware through a computer program, and the computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above-mentioned methods. Among them, any reference to memory, storage, database or other media used in the embodiments provided in this application may include non-volatile and / or volatile memory. Non-volatile memory may include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM) or flash memory. Volatile memory may include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM is available in many forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), double data rate SDRAM (DDRSDRAM), enhanced SDRAM (ESDRAM), synchronous link (Synchlink) DRAM (SLDRAM), memory bus (Rambus) direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and memory bus dynamic RAM (RDRAM), etc.
[0102] Those skilled in the art can clearly understand that for the convenience and simplicity of description, only the division of the above-mentioned functional units and modules is used as an example. In actual applications, the above-mentioned functions can be distributed and completed by different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above.
[0103] The embodiments described above are only used to illustrate the technical solutions of the present invention, rather than to limit the same. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that the technical solutions described in the aforementioned embodiments may still be modified, or some of the technical features may be replaced by equivalents. Such modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention, and should all be included in the protection scope of the present invention.
Claims
1. A classification model training method, characterized in that: The method comprises: Acquire training images, each of which is annotated with a category label and a position label of a detection target; Inputting the training image into a basic model to iteratively train the basic model, wherein the basic model includes a trunk classification network and a branch regression network, wherein the trunk classification network is used to perform target classification on the training image, and the branch regression network is used to obtain a feature image from the trunk classification network to perform target detection on the training image; After the basic model is trained, the branch regression network is discarded and the trunk classification network is determined as the target classification model.
2. The classification model training method according to claim 1, characterized in that: The step of inputting the training image into the basic model to iteratively train the trunk classification network and the branch regression network of the basic model comprises: Inputting the training image into the basic model to obtain the predicted classification result output by the trunk classification network and the first predicted coordinate output by the branch regression network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, a total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
3. The classification model training method according to claim 2, characterized in that: The branched regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
4. An image classification method, characterized in that: The method comprises: Get the target image; Inputting the target image into the target classification model to obtain a probability value that the target image belongs to each classification category, wherein the target classification model is trained by the classification model training method according to any one of claims 1 to 3; Determine that the classification category with the maximum probability value is the target category of the target image.
5. The image classification method according to claim 4, characterized in that: The acquiring of the target image comprises: Get the original image; Inputting the original image into a target detection model to obtain at least one second predicted coordinate and a prediction confidence of the second predicted coordinate; Determining whether the prediction confidence exceeds a preset confidence threshold; If the prediction confidence exceeds a preset confidence threshold, the target image is captured from the second prediction coordinates.
6. A classification model training device, characterized in that: include: A training image acquisition module, used to acquire training images, each of which is annotated with a category label and a position label of a detection target; A classification model training module, used for inputting the training image into a basic model to iteratively train the basic model, wherein the basic model includes a trunk classification network and a branch regression network, and the branch regression network obtains a feature image from the trunk classification network; The classification model determination module is used to discard the branch regression network and determine the trunk classification network as the target classification model after the basic model is trained.
7. The classification model training device according to claim 6, characterized in that: The classification model training module is used for: Inputting the training image into a basic model, obtaining a predicted classification result output by a trunk classification network of the basic model and a first predicted coordinate output by a branch regression network, wherein the branch regression network obtains a feature image from the trunk classification network; A classification loss value is calculated based on the predicted classification result and the category label, and a regression loss value is calculated based on the first predicted coordinates and the position label; Based on the classification loss value and the regression loss value, a total loss value is calculated; When the total loss value is less than a preset loss threshold, it is determined that the basic model is fully trained.
8. The classification model training device according to claim 7, characterized in that: The branched regression network outputs the predicted position result in the following manner: Acquire the feature image extracted by at least one feature extraction stage of the backbone classification network; Unifying the number of channels of each of the feature images and performing global pooling to obtain a standard feature image corresponding to each of the feature images; splicing the standard feature images to obtain a target feature image; The target feature image is input into a position prediction layer to obtain a first predicted coordinate of the training image.
9. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that: When the processor executes the computer program, the classification model training method as described in any one of claims 1 to 3 is implemented.
10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, it implements the classification model training method as described in any one of claims 1 to 3.