classify.cc

浏览该文件的文档.

/******************************************************************************
 * Copyright 2018 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the License);
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an AS IS BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
 
#include "modules/perception/traffic_light_recognition/recognition/caffe_recognizer/classify.h"
 
#include "cyber/common/file.h"
#include "cyber/profiler/profiler.h"
#include "modules/perception/common/camera/common/util.h"
#include "modules/perception/common/inference/inference_factory.h"
#include "modules/perception/common/inference/model_util.h"
#include "modules/perception/common/inference/utils/resize.h"
#include "modules/perception/common/util.h"
 
namespace apollo {
namespace perception {
namespace trafficlight {
 
void ClassifyBySimple::Init(const ClassifyParam& model_config,
                            const int gpu_id) {
  AINFO << "Enter Classify init";
  net_inputs_.clear();
  net_outputs_.clear();
  AINFO << "clear success";
  net_inputs_ =
      inference::GetBlobNames(model_config.info().inputs());
  net_outputs_ =
      inference::GetBlobNames(model_config.info().outputs());
 
  AINFO << net_inputs_.size();
  AINFO << net_outputs_.size();
 
  for (auto name : net_inputs_) {
    AINFO << "net input blobs: " << name;
  }
  for (auto name : net_outputs_) {
    AINFO << "net output blobs: " << name;
  }
 
  std::string model_path = GetModelPath(model_config.info().name());
 
  std::string proto_file =
      GetModelFile(model_path, model_config.info().proto_file().file());
  AINFO << "proto_file " << proto_file;
 
  std::string weight_file =
      GetModelFile(model_path, model_config.info().weight_file().file());
 
  AINFO << "weight_file" << weight_file;
 
  const auto& model_type = model_config.info().framework();
  rt_net_.reset(inference::CreateInferenceByName(model_type, proto_file,
                                                 weight_file, net_outputs_,
                                                 net_inputs_, model_path));
 
  rt_net_->set_gpu_id(gpu_id);
  gpu_id_ = gpu_id;
 
  resize_height_ = model_config.classify_resize_height();
  resize_width_ = model_config.classify_resize_width();
  unknown_threshold_ = model_config.classify_threshold();
  mean_.reset(new base::Blob<float>(1, 3, 1, 1));
  float* p = mean_->mutable_cpu_data();
  if (model_config.is_bgr()) {
    data_provider_image_option_.target_color = base::Color::BGR;
    p[0] = model_config.mean_b();
    p[1] = model_config.mean_g();
    p[2] = model_config.mean_r();
  } else {
    data_provider_image_option_.target_color = base::Color::RGB;
    p[0] = model_config.mean_r();
    p[1] = model_config.mean_g();
    p[2] = model_config.mean_b();
  }
  scale_ = model_config.scale();
 
  std::vector<int> shape = {1, resize_height_, resize_width_, 3};
  mean_buffer_.reset(new base::Blob<float>(shape));
 
  std::map<std::string, std::vector<int>> input_reshape;
  inference::AddShape(&input_reshape, model_config.info().inputs());
 
  AINFO << "input_reshape: " << input_reshape[net_inputs_[0]][0] << ", "
        << input_reshape[net_inputs_[0]][1] << ", "
        << input_reshape[net_inputs_[0]][2] << ", "
        << input_reshape[net_inputs_[0]][3];
 
  if (!rt_net_->Init(input_reshape)) {
    AWARN << "net init fail.";
  }
 
  image_.reset(
      new base::Image8U(resize_height_, resize_width_, base::Color::BGR));
}
 
void ClassifyBySimple::Perform(const camera::TrafficLightFrame* frame,
                               std::vector<base::TrafficLightPtr>* lights) {
  if (cudaSetDevice(gpu_id_) != cudaSuccess) {
    AERROR << "Failed to set device to " << gpu_id_;
    return;
  }
  std::shared_ptr<base::Blob<uint8_t>> rectified_blob;
 
  auto input_blob_recog = rt_net_->get_blob(net_inputs_[0]);
  auto output_blob_recog = rt_net_->get_blob(net_outputs_[0]);
 
  for (base::TrafficLightPtr light : *lights) {
    if (!light->region.is_detected) {
      continue;
    }
 
    data_provider_image_option_.crop_roi = light->region.detection_roi;
    data_provider_image_option_.do_crop = true;
    data_provider_image_option_.target_color = base::Color::BGR;
    frame->data_provider->GetImage(data_provider_image_option_, image_.get());
 
    AINFO << "get img done";
 
    const float* mean = mean_.get()->cpu_data();
    inference::ResizeGPU(*image_, input_blob_recog,
                         frame->data_provider->src_width(), 0, mean[0], mean[1],
                         mean[2], true, scale_);
 
    AINFO << "resize gpu finish.";
    cudaDeviceSynchronize();
 
    PERF_BLOCK("traffic_light_recognition_inference")
    rt_net_->Infer();
    PERF_BLOCK_END
 
    cudaDeviceSynchronize();
    AINFO << "infer finish.";
 
    float* out_put_data = output_blob_recog->mutable_cpu_data();
    Prob2Color(out_put_data, unknown_threshold_, light);
  }
}
 
void ClassifyBySimple::Prob2Color(const float* out_put_data, float threshold,
                                  base::TrafficLightPtr light) {
  int max_color_id = 0;
  std::vector<base::TLColor> status_map = {
      base::TLColor::TL_BLACK, base::TLColor::TL_RED, base::TLColor::TL_YELLOW,
      base::TLColor::TL_GREEN};
  std::vector<std::string> name_map = {"Black", "Red", "Yellow", "Green"};
  std::vector<float> prob(out_put_data, out_put_data + status_map.size());
  auto max_prob = std::max_element(prob.begin(), prob.end());
  max_color_id = (*max_prob > threshold)
                     ? static_cast<int>(std::distance(prob.begin(), max_prob))
                     : 0;
 
  light->status.color = status_map[max_color_id];
  light->status.confidence = out_put_data[max_color_id];
  AINFO << "Light status recognized as " << name_map[max_color_id];
  AINFO << "Color Prob:";
  for (size_t j = 0; j < status_map.size(); j++) {
    AINFO << out_put_data[j];
  }
}
 
}  // namespace trafficlight
}  // namespace perception
}  // namespace apollo