2017년 6월 3일 토요일

opencv Mat 를 caffe datum으로 변환, datum을 Mat로 변환

 기존 opencv에서 영상처리 하던게 익숙해서 caffe를 Mat로 변환 후에 Opencv 로 처리 할 경우가 많은데, bvlc caffe에 아직 float형 image와 multi channel(4개 이상) 이미지 변환함수가 안 만들어져 있어서 만들어 보았습니다.

 아래는 Mat를 datum으로 변환하는 것.
void CVMatToDatum(const cv::Mat& cv_img, Datum* datum, int depth)
{
 if (depth == CV_8U)
 {
  CVMatToDatum(cv_img, datum);
  return;
 }

 CHECK(cv_img.depth() == CV_32F) << "data type must be 32bit float";

 datum->set_channels(cv_img.channels());
 datum->set_height(cv_img.rows);
 datum->set_width(cv_img.cols);
 datum->clear_data();
 datum->clear_float_data();
 datum->set_encoded(false);
 int datum_channels = datum->channels();
 int datum_height = datum->height();
 int datum_width = datum->width();
 int datum_size = datum_channels * datum_height * datum_width;

 for (int h = 0; h < datum_height; ++h) 
 {
  const float* ptr = cv_img.ptr<float>(h);
  int img_index = 0;
  for (int w = 0; w < datum_width; ++w) {
   for (int c = 0; c < datum_channels; ++c) {
    datum->add_float_data( ptr[img_index++]);
   }
  }
 }
 return;
}

void CVMatToDatum(const cv::Mat& cv_img, Datum* datum) {
  CHECK(cv_img.depth() == CV_8U) << "Image data type must be unsigned byte";
  datum->set_channels(cv_img.channels());
  datum->set_height(cv_img.rows);
  datum->set_width(cv_img.cols);
  datum->clear_data();
  datum->clear_float_data();
  datum->set_encoded(false);
  int datum_channels = datum->channels();
  int datum_height = datum->height();
  int datum_width = datum->width();
  int datum_size = datum_channels * datum_height * datum_width;
  std::string buffer(datum_size, ' ');
  for (int h = 0; h < datum_height; ++h) {
    const uchar* ptr = cv_img.ptr<uchar>(h);
    int img_index = 0;
    for (int w = 0; w < datum_width; ++w) {
      for (int c = 0; c < datum_channels; ++c) {
        int datum_index = (c * datum_height + h) * datum_width + w;
        buffer[datum_index] = static_cast<char>(ptr[img_index++]);
      }
    }
  }
  datum->set_data(buffer);
}

 아래는 datum을 Mat로 변환하는 것.
cv::Mat DatumToCVMat(const Datum& datum)
{
 int datum_channels = datum.channels();
 int datum_height = datum.height();
 int datum_width = datum.width();

 string strData = datum.data();
 cv::Mat cv_img;
 if (strData.size() != 0)
 {
  cv_img.create(datum_height, datum_width, CV_8UC(datum_channels));
  const string& data = datum.data();
  std::vector<char> vec_data(data.c_str(), data.c_str() + data.size());

  for (int h = 0; h < datum_height; ++h) {
   uchar* ptr = cv_img.ptr<uchar>(h);
   int img_index = 0;
   for (int w = 0; w < datum_width; ++w) {
    for (int c = 0; c < datum_channels; ++c) {
     int datum_index = (c * datum_height + h) * datum_width + w;
     ptr[img_index++] = static_cast<uchar>(vec_data[datum_index]);
    }
   }
  }
 }
 else
 {
  cv_img.create(datum_height, datum_width, CV_32FC(datum_channels));
  for (int h = 0; h < datum_height; ++h) {
   float* ptr = cv_img.ptr<float>(h);
   int img_index = 0;
   for (int w = 0; w < datum_width; ++w) {
    for (int c = 0; c < datum_channels; ++c) {
     ptr[img_index++] = static_cast<float>(datum.float_data(img_index));
    }
   }
  }
 }
 return cv_img;
}

댓글 없음:

댓글 쓰기