Vector magnetometer is an important branch of quantum precision measurement, which has been widely used in basic physics, biomedicine, material science and other fields. As a quantum sensor, the Nitrogen-Vacancy Center has a longer coherency time at room temperature, allowing for nanoscale magnetic field detection. Three different coaxial NV centers of diamond should be selected as magnetic field sensors to realize the reconstruction of vector magnetic field withnanometer resolution. Whether the axial information of NV center can be obtained quickly and accurately directly affects the accuracy and efficiency of vector magnetic field measurement. Therefore, how to obtain the axial information of the NV center is influential. This paper proposes a confocal microscopy imaging system based on machine learning fitting method, combined with an angularly polarized beam instead of a Gaussian beam to excite the NV center to obtain a scanning fluorescence image. Then the identification of the diamond NV center axis and the problem of calculating the azimuth and polar angle of the NV axis in space are solved through the method. This article demonstrates a model based on convolutional neural network to automatically identify and locate NV centers in diamond, and on the basis of object detection, the method of gradient descent fitting is combined with image recognition processing and matching algorithms. The model accelerates extraction of the axial information of the four NV axes, optimizes the extraction process of the NV center axial information, improves the speed and accuracy of fitting, and thereby increases the efficiency of magnetic field vector reconstruction. The research results of this design can be applied to the measurement of the magnetic field vector, and not only can it improve the calculation efficiency of the NV center angle direction in the diamond, but also has a degree of robustness to the noise in the experimental environment.