Hu Chunming / village_ascend_arm

#include "PicAnalysis.h"
#include "./utils/logger.hpp"


PicAnalysis::PicAnalysis(/* args */)
{
    aclInit(nullptr);
}

PicAnalysis::~PicAnalysis()
{
    LOG_INFO("~PicAnalysis()");
    release();
    aclFinalize();
}

int PicAnalysis::init(VillageParam param) {

    string log_path = param.log_path + "/main.log";

    set_default_logger(LogLevel(param.log_level), "PicAnalysis", log_path.c_str(), param.log_mem, param.log_days);

    // 输入参数
    LOG_INFO("dev_id: {}", param.dev_id);
    LOG_INFO("sdk_path: {}", param.sdk_path);
    LOG_INFO("log_level: {}", param.log_level);
    LOG_INFO("log_days: {}", param.log_days);
    LOG_INFO("log_path: {}", log_path);

    int dev_id = param.dev_id;

    int ret = SY_FAILED;

    ret = m_vehicle_analysis.init(dev_id, param.sdk_path, 16);
    if(0 != ret){
        return -1;
    }

    head_tail_param ht_param;
    ht_param.devId = dev_id;
    ht_param.max_batch = 16;
    ht_param.sdk_root = param.sdk_path;
    ret = m_head_tail_algorithm.init(ht_param);
    if(0 != ret){
        return -1;
    }

    ret = m_clothes_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_human_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_human_car_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_motor_rainshed_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_motor_phone_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_road_seg_algorithm.init(dev_id, param.sdk_path);
    if(0 != ret){
        return -1;
    }

    ret = m_crop_util.init(dev_id);
    if(0 != ret){
        return -1;
    }

    ACL_CALL(aclrtCreateContext(&m_ctx, 0), ACL_SUCCESS, SY_FAILED);
	ACL_CALL(aclrtSetCurrentContext(m_ctx), ACL_SUCCESS, SY_FAILED);
	ACL_CALL(aclrtCreateStream(&stream), ACL_SUCCESS, SY_FAILED);
    m_dvpp = new DvppProcess();
	m_dvpp->InitResource(stream);

    return 0;
}

vector<AnalysisResult> PicAnalysis::analysis_file(vector<string> vec_file_path){

    vector<AnalysisResult> result;

    int ret = aclrtSetCurrentContext(m_ctx);
    if (SY_SUCCESS != ret) {
        printf("aclrtSetCurrentContext failed!");
        return result;
    }

    const int batch_size = vec_file_path.size();

    vector<sy_img> vec_img;

    // ImageData 内部是智能指针，分析未处理完成前不得释放
    ImageData dvpp_data[batch_size];
    for (size_t i = 0; i < vec_file_path.size(); i++)
    {
        string file_path = vec_file_path[i];
        ImageData src;
        ret = Utils::ReadImageFile(src, file_path); //将二进制图像读入内存，并读取宽高信息
        if(ret != SY_SUCCESS){
            LOG_ERROR("ReadImageFile failed!");
            return result;
        }

        ret = m_dvpp->CvtJpegToYuv420sp(dvpp_data[i], src); //解码
        if(ret != SY_SUCCESS){
            LOG_ERROR("CvtJpegToYuv420sp failed!");
            return result;
        }
        
        sy_img img;
        img.w_ = dvpp_data[i].width;
        img.h_ = dvpp_data[i].height;
        img.data_ = dvpp_data[i].data.get();

        vec_img.push_back(img);
    }

    return analysis_img(vec_img);
}

vector<AnalysisResult> PicAnalysis::va_result2AnalysisResult(va_result* result, int batchsize) {

    vector<AnalysisResult> vec_result;

    for (int b = 0; b < batchsize; b++)
    {
        vector<VehicleInfo> vec_info;
        for(int c=0;c<result[b].count;c++)
        {
            vehicle_info result_info = result[b].info[c];

            VehicleInfo info;
            info.vehicle_detect_res         = result_info.vehicle_detect_res        ;	 
            info.vehicle_win_detect_res     = result_info.vehicle_win_detect_res    ; 	
            info.vehicle_body_detect_res    = result_info.vehicle_body_detect_res   ;	
            info.vehicle_color_res          = result_info.vehicle_color_res         ;
            info.vehicle_recg_res           = result_info.vehicle_recg_res          ;  
            info.vehicle_plate_det_recg_res = result_info.vehicle_plate_det_recg_res;	 
            info.vehicle_illegal_det_res    = result_info.vehicle_illegal_det_res   ; 	
            info.vehicle_fea_res            = result_info.vehicle_fea_res           ;   	
            info.mta_res                    = result_info.mta_res                   ;   
            info.manned_res	 	            = result_info.manned_res	 	        ;  
            info.type                       = result_info.type                      ;         
            info.vpt_type                   = result_info.vpt_type                  ;


            auto pendant_result = result_info.vehicle_pendant_det_res;
            for (size_t j = 0; j < pendant_result.count; j++)
            {
                auto pendant_det_info = pendant_result.vpd_res[j];

                pendant_info one_pendant_info;
                one_pendant_info.rect = pendant_det_info.rect;
                one_pendant_info.index = pendant_det_info.index;
                one_pendant_info.confidence = pendant_det_info.confidence;
                one_pendant_info.driver_copilot_info = pendant_det_info.driver_copilot_info;
                info.vehicle_pendant_det_res.push_back(one_pendant_info);
            }

            vec_info.push_back(info);
        }

        AnalysisResult one_result;
        one_result.info = vec_info;
        vec_result.push_back(one_result);
    }

    return vec_result;
}

vector<AnalysisResult> PicAnalysis::analysis_img(vector<sy_img> vec_img){

    LOG_INFO("vec_img size:{}", vec_img.size());

    vector<AnalysisResult> vec_result;

    const int batch_size = vec_img.size();

    int ret = aclrtSetCurrentContext(m_ctx);
    if (SY_SUCCESS != ret) {
        LOG_INFO("aclrtSetCurrentContext failed!");
        return vec_result;
    }

    va_result* result = m_vehicle_analysis.detect(vec_img);
    if (result) {
        vec_result = va_result2AnalysisResult(result, batch_size);
        m_vehicle_analysis.release_result(result, vec_img.size());
    }

    LOG_INFO("vec_result size:{}", vec_result.size());

    bool bConsistent = false;
    std::vector<RoadInfo> vec_road = m_road_seg_algorithm.detect(vec_img);
    if (vec_road.size() == batch_size && vec_result.size() == batch_size) {
        bConsistent = true;
        for (size_t i = 0; i < batch_size; i++) {
            vec_result[i].vec_line = vec_road[i].vec_line;
            vec_result[i].vec_road = vec_road[i].vec_road;
        }
    }

    for (int b = 0; b < vec_result.size(); b++)
    {
        vector<VehicleInfo>& vec_info = vec_result[b].info;

        RoadInfo road_info;
        if (bConsistent) {
            road_info = vec_road[b];
        }

        sy_img img = vec_img[b];

        ImageData src;
        src.width = img.w_;
        src.height = img.h_;
        src.alignWidth = ALIGN_UP16(img.w_);
        src.alignHeight = ALIGN_UP2(img.h_);
        src.size = YUV420SP_SIZE(src.alignWidth, src.alignHeight);
        src.data_naked = img.data_;
        for(int c=0;c<vec_info.size();c++)
        {
            VehicleInfo& result_info = vec_info[c];
            int shot_type=result_info.type;

            sy_rect vehicle_rect = result_info.vehicle_body_detect_res.rect;

            if (road_info.vec_direct.size() > 0 && (1 == shot_type || 0 == shot_type)) {
                result_info.reverse_driving = m_road_seg_algorithm.check_reverse_driving(road_info.vec_direct, vehicle_rect, src.width, src.height, shot_type);
            }

            // 行人
            if (0 == result_info.vpt_type)
            {
                ImageData* human_data = m_crop_util.crop(src, vehicle_rect.left_, vehicle_rect.top_, vehicle_rect.left_ + vehicle_rect.width_, vehicle_rect.top_ + vehicle_rect.height_);

                sy_img img;
                img.w_ = human_data->alignWidth;
                img.h_ = human_data->alignHeight;
                img.data_ = human_data->data_naked;

                vector<sy_img> vec_human_img;
                vec_human_img.push_back(img);

                vector<BodyColorInfo> vec_body_color = m_human_algorithm.detect(vec_human_img);

                result_info.human_upper_color = vec_body_color[0].upper_body_color;
                result_info.human_lower_color = vec_body_color[0].lower_body_color;

                delete human_data;
                human_data = nullptr;
            } else if(result_info.vpt_type == 1 || result_info.vpt_type == 2 || result_info.vpt_type == 3){
                if(result_info.manned_res.hs_count >= 3){
                    // 摩托车、三轮车载人
                    result_info.motor_manned = 1;
                }	
            } else {
                ImageData* vehicle_data = m_crop_util.crop(src, vehicle_rect.left_, vehicle_rect.top_, vehicle_rect.left_ + vehicle_rect.width_, vehicle_rect.top_ + vehicle_rect.height_);

                sy_img img;
                img.w_ = vehicle_data->alignWidth;
                img.h_ = vehicle_data->alignHeight;
                img.data_ = vehicle_data->data_naked;

                vector<sy_img> vec_vehicle_img;
                vec_vehicle_img.push_back(img);

                vector<HeadTailResult> head_tail_result;
                ret = m_head_tail_algorithm.detect(vec_vehicle_img, head_tail_result);
                if (0 != ret) {
                    LOG_ERROR("m_head_tail_algorithm failed!");
                    head_tail_result.clear();
                } else {
                    auto one_ht_result = head_tail_result[0];

                    if(result_info.vpt_type == 6 || result_info.vpt_type == 7){
                        if (one_ht_result.cls == 0 && result_info.manned_res.hs_count > 0)
                        { // 车尾，判断是否 货车尾部货厢载人
                            result_info.truck_manned = 1;
                        }	
                    } 
                }
                delete vehicle_data;
                vehicle_data = nullptr;
            }

            // 司乘
            auto& pendant_res = result_info.vehicle_pendant_det_res;
            int vpd_num = pendant_res.size();
            vector<sy_img> vec_human_img;
            vector<ImageData*> vec_data;
            for(int p=0; p<vpd_num; p++)
            {
                int index = pendant_res[p].index;
                if(index == 0){ //driver
                    sy_rect human_rect = pendant_res[p].rect;
                    ImageData* human_data = m_crop_util.crop(src, human_rect.left_, human_rect.top_, human_rect.left_ + human_rect.width_, human_rect.top_ + human_rect.height_);

                    sy_img img;
                    img.w_ = human_data->alignWidth;
                    img.h_ = human_data->alignHeight;
                    img.data_ = human_data->data_naked;

                    vec_human_img.push_back(img);

                    vec_data.push_back(human_data);
                }
            }

            vector<int> vec_color = m_clothes_algorithm.detect(vec_human_img);

            for(int p=0; p<vec_color.size(); p++)
            {
                int index = pendant_res[p].index;
                if(index == 0) { // 更新司乘衣着颜色
                    pendant_res[p].iColor = vec_color[p];

                    ImageData* human_data = vec_data[p];
                    delete human_data;
                    human_data = nullptr;
                }
            }

            // 摩托车
            if(shot_type==2)//摩托车
            {
                ImageData* motor_data = m_crop_util.crop(src, vehicle_rect.left_, vehicle_rect.top_, vehicle_rect.left_ + vehicle_rect.width_, vehicle_rect.top_ + vehicle_rect.height_);

                sy_img img;
                img.w_ = motor_data->alignWidth;
                img.h_ = motor_data->alignHeight;
                img.data_ = motor_data->data_naked;

                vector<sy_img> vec_motor_img;
                vec_motor_img.push_back(img);

                vector<MotorRainshedResult> vec_rainshed_result = m_motor_rainshed_algorithm.detect(vec_motor_img);
                if (vec_rainshed_result.size() > 0) {
                    result_info.rainshed = vec_rainshed_result[0].rainshed;
                    result_info.rainshed_prob = vec_rainshed_result[0].prob;
                }

                vector<MotorPhoneResult> vec_phone_result = m_motor_phone_algorithm.detect(vec_motor_img);
                if (vec_phone_result.size() > 0) {
                    result_info.phoning = vec_phone_result[0].phoning;
                    result_info.phoning_prob = vec_phone_result[0].prob;
                }

                int iColor = -1;
                std::vector<HumanCarResult> vec_hcp_result = m_human_car_algorithm.detect(vec_motor_img);
                if (vec_hcp_result.size() > 0 && road_info.vec_direct.size() > 0) {
                    int head_or_tail = vec_hcp_result[0].orient;
                    if (head_or_tail == 0 || head_or_tail == 1) {
                        result_info.reverse_driving = m_road_seg_algorithm.check_reverse_driving(road_info.vec_direct, vehicle_rect, src.width, src.height, head_or_tail);
                    }
                    iColor = vec_hcp_result[0].up_color;
                }

                int hs_num = result_info.manned_res.hs_count;
                for (size_t i = 0; i < hs_num; i++)
                {
                    pendant_info one_pendant;
                    one_pendant.confidence = result_info.manned_res.hs_rect[i].score;
                    one_pendant.rect = result_info.manned_res.hs_rect[i].rect;
                    one_pendant.index = 0;
                    one_pendant.iColor = iColor;

                    result_info.vehicle_pendant_det_res.push_back(one_pendant);
                    
                }

                delete motor_data;
                motor_data = nullptr;
            }

            //压黄实线
            result_info.cross_line = m_road_seg_algorithm.check_cross_line(road_info.vec_line, vehicle_rect, src.width, src.height);
            // 压导流线
            result_info.cross_diversion_line = m_road_seg_algorithm.check_cross_region(road_info.vec_road, vehicle_rect, src.width, src.height, 3);    //3是导流线区域
        }
    }

    LOG_INFO("analysis finished!");

    return vec_result;
}

int PicAnalysis::release() {

    ACL_CALL(aclrtSetCurrentContext(m_ctx), ACL_SUCCESS, SY_FAILED);

    delete m_dvpp;
    m_dvpp = nullptr;

    if (stream != nullptr) {
        int ret = aclrtDestroyStream(stream);
        if (ret != ACL_SUCCESS) {
            LOG_ERROR("destroy stream failed");
        }
        stream = nullptr;
    }

    aclrtDestroyContext(m_ctx);

    return 0;
}