src/ai_engine_module/VehicleAnalysis.cpp

#include "VehicleAnalysis.h"
VehicleAnalysis::VehicleAnalysis() {
    cout << va_get_version() << endl;
}
VehicleAnalysis::~VehicleAnalysis() {
    release();
}
int VehicleAnalysis::init(int devId, std::string sdk_root, int max_batch_size) {
    
	param.vehicle_detect_config= SY_CONFIG_OPEN;  //1.开启车检测 SY_CONFIG_CLOSE  SY_CONFIG_OPEN
    param.vehicle_recg_config= SY_CONFIG_OPEN;     //4.开启车型识别
	param.vehicle_recg_supplement_config= SY_CONFIG_OPEN;     //4.开启车型识别补充识别
    param.vehicle_plate_det_recg_config= SY_CONFIG_OPEN; //5.开启车牌检测识别	
    param.vehicle_pendant_det_config= SY_CONFIG_OPEN; //6.开启车属性检测识别
    param.vehicle_motor_tricycle_analysis_config= SY_CONFIG_OPEN;  //8.摩托车分析	
    param.vehicle_manned_config= SY_CONFIG_OPEN;  //8.开启载人分析	
    param.vehicle_illegal_config= SY_CONFIG_OPEN;  //7.开启车违规
	param.vehicle_color_config= SY_CONFIG_CLOSE;  //3.开启车颜色识别
	param.vehicle_feature_config= SY_CONFIG_CLOSE;  //8.开启车辆特征提取
	
	param.gpuid=devId;
		
	//车检测参数	
	//param.vehicle_det_param.process_min_l = 720;// 720;
	//param.vehicle_det_param.process_max_l = 1280;//1280;	
	//param.vehicle_det_param.thresld=0.3;
	//param.min_obj_size=200;  //最小目标	
	param.vehicle_det_thresld=0.4;
	
	//车牌检测参数
	//param.vehicle_plate_det_param.process_min_l=320;
	//param.vehicle_plate_det_param.process_max_l=320;	
	//param.vehicle_plate_det_param.thresld=0.4;
	param.vehicle_plate_det_thresld=0.5;
	
	//车属性检测参数
	//param.vehicle_attribute_det_param.process_min_l=360;
	//param.vehicle_attribute_det_param.process_max_l=640;	
	//param.vehicle_attribute_det_param.thresld=0.3;
	param.vehicle_attribute_det_thresld=0.5;
	
	//车logo检测参数
	//param.vehicle_logo_det_param.process_min_l=512;
	//param.vehicle_logo_det_param.process_max_l=512;	
	//param.vehicle_logo_det_param.thresld=0.1;
	param.vehicle_logo_det_thresld=0.1;	
	
	//车颜色阈值
	param.vc_thresld = 0.5;
	
    string dbPath = sdk_root + "/models/vehicle_analysis/db/vr_h0725x210605_r191230.db";
    string models_path = sdk_root + "/models/vehicle_analysis/";
	//车型参数
	param.dbPath= (char*)dbPath.data();	
	param.models_Path= (char*)models_path.data();	//所有模型的地址
    LOG_INFO("va_init start");
    // 内部有 ctx
    int ret = va_init(&m_handle, param);
    if (ret != 0) {
        return -1;
	}
    LOG_INFO("va_init success");
    return 0;
}
va_result * VehicleAnalysis::detect(vector<sy_img> vec_img) {
			
    int batch_size = vec_img.size();
    
    va_result *result=new va_result[batch_size];
    for(int b=0;b<batch_size;b++)
    {
        result[b].count=0;
        result[b].info=new vehicle_info[100];
        for(int c=0; c<100; c++)
        {
            result[b].info[c].vehicle_pendant_det_res.vpd_res=new v_pendant_d_info[300];	
        }			
    }
    
    int ret = va_batch(m_handle, vec_img.data(), batch_size, result);
    vector<VehicleAnalysisResult> vec_result;
    
    for (int b = 0; b < batch_size; b++)
    {						
        for(int c=0;c<result[b].count;c++)
        {
            vehicle_info result_info = result[b].info[c];
            VehicleAnalysisResult analysis_result;
            std::string str_vehicle_type;
            int shot_type=result_info.type;					
            if (shot_type==0)str_vehicle_type = "head";
            else if (shot_type==1)str_vehicle_type = "rear";
            else if (shot_type==2)str_vehicle_type = "motor";					
            else if (shot_type==3)str_vehicle_type = "tricycle";
            else if (shot_type==4)str_vehicle_type = "body";
            else if (shot_type==5)str_vehicle_type = "body_nova";
            
            analysis_result.shot_type = shot_type;
            
           //1.车头结果：蓝框//2.车窗结果：绿框
            if(param.vehicle_detect_config==SY_CONFIG_OPEN)
            {
                //车身
                //std::cout << "vehicle_body_detect_res info:"<< std::endl;
                if(result_info.vehicle_body_detect_res.rect.width_>0)
                {
                    float vehicle_body_detect_res_score = result_info.vehicle_body_detect_res.score;
                    int x1=result_info.vehicle_body_detect_res.rect.left_;
                    int y1=result_info.vehicle_body_detect_res.rect.top_;
                    int x2=result_info.vehicle_body_detect_res.rect.left_+result_info.vehicle_body_detect_res.rect.width_;
                    int y2=result_info.vehicle_body_detect_res.rect.top_+result_info.vehicle_body_detect_res.rect.height_;					
                    std::cout << "		    vehicle_body_detect_res_score:" <<vehicle_body_detect_res_score<<",car_rect:[" <<x1<<"," <<y1<<"," <<x2-x1<<"," <<y2-y1<<"]"<< std::endl;
                        
                    analysis_result.vehicle_rect = result_info.vehicle_body_detect_res.rect;				
                }	
                
                // //车头
                // if(result_info.vehicle_detect_res.rect.width_>0)	
                // {						
                //     float vehicle_detect_res_score = result_info.vehicle_detect_res.score;
                //     int x1=result_info.vehicle_detect_res.rect.left_;
                //     int y1=result_info.vehicle_detect_res.rect.top_;
                //     int x2=result_info.vehicle_detect_res.rect.left_+result_info.vehicle_detect_res.rect.width_;
                //     int y2=result_info.vehicle_detect_res.rect.top_+result_info.vehicle_detect_res.rect.height_;										
                //     std::cout << "		    vehicle_detect_res_score:" <<vehicle_detect_res_score<<",car_rect:[" <<x1<<"," <<y1<<"," <<x2-x1<<"," <<y2-y1<<"]"<< std::endl;
                // }
                
                // //车窗
                // if(result_info.vehicle_win_detect_res.rect.width_>0)
                // {
                //     float vehicle_win_detect_res_score = result_info.vehicle_win_detect_res.score;
                //     int x1=result_info.vehicle_win_detect_res.rect.left_;
                //     int y1=result_info.vehicle_win_detect_res.rect.top_;
                //     int x2=result_info.vehicle_win_detect_res.rect.left_+result_info.vehicle_win_detect_res.rect.width_;
                //     int y2=result_info.vehicle_win_detect_res.rect.top_+result_info.vehicle_win_detect_res.rect.height_;					
                //     std::cout << "		    vehicle_win_detect_res_score:" <<vehicle_win_detect_res_score<<",win_rect:[" <<x1<<"," <<y1<<"," <<x2-x1<<"," <<y2-y1<<"]"<< std::endl;
                // }					
            }
            //4.VR车型识别
            if(param.vehicle_recg_config==SY_CONFIG_OPEN  && (shot_type==0 || shot_type==1 || shot_type==4 ))
            {
                char *vehicle_brand=result_info.vehicle_recg_res.vehicle_brand;//车辆品牌
                char* vehicle_subbrand=result_info.vehicle_recg_res.vehicle_subbrand;             //车辆子品牌
                char* vehicle_issue_year=result_info.vehicle_recg_res.vehicle_issue_year;           //车辆年款
                char* vehicle_type_=result_info.vehicle_recg_res.vehicle_type;                //车辆类型
                char* freight_ton=result_info.vehicle_recg_res.freight_ton;                  //货车吨级
                float name_score=result_info.vehicle_recg_res.name_score;                   //识别置信度
                
                float name_score_thre = 0;
                //if(shot_type==0)name_score_thre=0.7;//车头车型识别建议阈值0.7
                //if(shot_type==1)name_score_thre=0.8;//车尾车型识别建议阈值0.8
                if(name_score > name_score_thre)
                {
                    analysis_result.vehicle_type = vehicle_type_;
                }
            } else {
                analysis_result.vehicle_type = str_vehicle_type;
            }
            std::cout << "vehicle_type_: " << analysis_result.vehicle_type << std::endl;
            //5.VP车牌检测识别
            if(param.vehicle_plate_det_recg_config==SY_CONFIG_OPEN  && (shot_type==0 || shot_type==1))
            {
                vplate_results plate_result = result_info.vehicle_plate_det_recg_res;
                std::cout << "		car plate info:"<< endl;
                int special_type=plate_result.special_type;//常规车牌、临时车牌、低速车牌。0-common,1-temporary_license_plate,2-low_speed_license_plate.
                std::cout << "		    special_type:" << special_type<< std::endl;
                
                float detect_score=plate_result.detect_score;
                if(detect_score>0.3)
                {															
                    //车牌识别结果
                    std::string plate_recg="";
                    std::string character_prob;
                    //if(num_score>0.99)//车牌识别建议置信度阈值0.99
                    {
                        
                        for (int m = 0; m < PLATENUM; m++)
                        {
                            plate_recg=plate_recg + plate_result.recg[m].character;
                            character_prob += std::string(plate_result.recg[m].character) + "-" + std::to_string(plate_result.recg[m].maxprob) + ",";
                        }
                    }
                    analysis_result.plate_det_score = plate_result.detect_score;
                    analysis_result.plate_rect = plate_result.rect;//车牌检测坐标
                    analysis_result.plate_type = plate_result.type;	//车牌类型：0-单排蓝色 1-单排黄色 2-单排白色 3-单排黑色 4-双排黄色 5-双排白色 6-新能源黄绿色 7-新能源白绿色
                    analysis_result.plate_num = plate_recg;
                    analysis_result.character_prob = character_prob;
                    analysis_result.plate_num_score = plate_result.num_score;//识别置信度
                    // analysis_result.special_type;						//常规车牌、临时车牌、低速车牌。0-common,1-temporary_license_plate,2-low_speed_license_plate.
                }
                
            }
            //6.车属性结果：
            if(param.vehicle_pendant_det_config==SY_CONFIG_OPEN && (shot_type==0|| shot_type==4))
            {
                std::cout << "		vehicle_pendant_det_res info:"<< endl;
                int vpd_num = result_info.vehicle_pendant_det_res.count;
                //std::cout << vpd_num<< std::endl;
                std::cout << "		    vpd_num:"<<vpd_num<< endl;
                for(int p=0; p<vpd_num; p++)
                {
                    int index = result_info.vehicle_pendant_det_res.vpd_res[p].index;
                    if(index == 0){
                        DriverInfo info;
                        info.driver_rect = result_info.vehicle_pendant_det_res.vpd_res[p].rect;
                        info.driver_prob = result_info.vehicle_pendant_det_res.vpd_res[p].confidence;
                        analysis_result.vec_drivers.push_back(info);
                    }
                }
            }
            //11.摩托车三轮车分析
            if(param.vehicle_motor_tricycle_analysis_config==SY_CONFIG_OPEN && (shot_type==2 || shot_type==3))
            {
                //std::cout << "mta output----"<< std::endl;
                
                if(shot_type==2)//摩托车
                {
                    //摩托车载人
                    //int status=result_info.mta_res.motor_manned.status;	
                    //float score=result_info.mta_res.motor_manned.confidence;
                    //if(status == MOTOR_MANNED)
                    //	std::cout << "motor info: MOTOR_MANNED, prob: " << score << std::endl;
                    //else if(status == MOTOR_NOT_MANNED)
                    //	std::cout << "motor info: MOTOR_NOT_MANNED, prob: " << score << std::endl;
                    
                    //摩托车驾驶人是否戴安全帽
                    analysis_result.motor_helmeted = result_info.mta_res.motor_driver_helmeted.status;
                    float score=result_info.mta_res.motor_driver_helmeted.confidence;
                
                }
            }
            //载人输出
            if(param.vehicle_manned_config==SY_CONFIG_OPEN && shot_type!=5)
            {
                if(result_info.manned_res.hs_count >= 3){
                    analysis_result.motor_manned = 1;
                }			
            }				
        }                            
    }
    
    //delete result
    // for(int b=0;b<batch_size;b++)
    // {
    //     for(int c=0; c<100; c++)
    //     {
    //         if(result[b].info[c].vehicle_pendant_det_res.vpd_res!=NULL)
    //             delete[] result[b].info[c].vehicle_pendant_det_res.vpd_res;
            
    //     }
    //     if(result[b].info!=NULL){
    //         delete[] result[b].info;
    //     }
            
    // }
    // if(result!=NULL){
    //     delete[] result;
    // }
    return result;  
}
void VehicleAnalysis::release(){
    if (m_handle) {
        va_release(&m_handle);
    }
}
void VehicleAnalysis::release_result(va_result* result, int batch_size){
    for(int b=0;b<batch_size;b++)
    {
        for(int c=0; c<100; c++)
        {
            if(result[b].info[c].vehicle_pendant_det_res.vpd_res!=NULL)
                delete[] result[b].info[c].vehicle_pendant_det_res.vpd_res;
            
        }
        if(result[b].info!=NULL){
            delete[] result[b].info;
        }
            
    }
    if(result!=NULL){
        delete[] result;
    }
}