misc update
[jacinto-ai/caffe-jacinto.git] / src / caffe / util / io.cpp
1 // Copyright 2013 Yangqing Jia
3 #include <stdint.h>
4 #include <fcntl.h>
5 #include <google/protobuf/text_format.h>
6 #include <google/protobuf/io/zero_copy_stream_impl.h>
7 #include <opencv2/core/core.hpp>
8 #include <opencv2/highgui/highgui.hpp>
10 #include <algorithm>
11 #include <string>
12 #include <iostream>
13 #include <fstream>
15 #include "caffe/common.hpp"
16 #include "caffe/util/io.hpp"
17 #include "caffe/proto/caffe.pb.h"
19 using cv::Mat;
20 using cv::Vec3b;
21 using std::fstream;
22 using std::ios;
23 using std::max;
24 using std::string;
25 using google::protobuf::io::FileInputStream;
26 using google::protobuf::io::FileOutputStream;
28 namespace caffe {
30 void ReadImageToProto(const string& filename, BlobProto* proto) {
31   Mat cv_img;
32   cv_img = cv::imread(filename, CV_LOAD_IMAGE_COLOR);
33   CHECK(cv_img.data) << "Could not open or find the image.";
34   DCHECK_EQ(cv_img.channels(), 3);
35   proto->set_num(1);
36   proto->set_channels(3);
37   proto->set_height(cv_img.rows);
38   proto->set_width(cv_img.cols);
39   proto->clear_data();
40   proto->clear_diff();
41   for (int c = 0; c < 3; ++c) {
42     for (int h = 0; h < cv_img.rows; ++h) {
43       for (int w = 0; w < cv_img.cols; ++w) {
44         proto->add_data(static_cast<float>(cv_img.at<Vec3b>(h, w)[c]) / 255.);
45       }
46     }
47   }
48 }
50 void ReadImageToDatum(const string& filename, const int label, Datum* datum) {
51   Mat cv_img;
52   cv_img = cv::imread(filename, CV_LOAD_IMAGE_COLOR);
53   CHECK(cv_img.data) << "Could not open or find the image.";
54   DCHECK_EQ(cv_img.channels(), 3);
55   datum->set_channels(3);
56   datum->set_height(cv_img.rows);
57   datum->set_width(cv_img.cols);
58   datum->set_label(label);
59   datum->clear_data();
60   datum->clear_float_data();
61   string* datum_string = datum->mutable_data();
62   for (int c = 0; c < 3; ++c) {
63     for (int h = 0; h < cv_img.rows; ++h) {
64       for (int w = 0; w < cv_img.cols; ++w) {
65         datum_string->push_back(static_cast<char>(cv_img.at<Vec3b>(h, w)[c]));
66       }
67     }
68   }
69 }
72 void WriteProtoToImage(const string& filename, const BlobProto& proto) {
73   CHECK_EQ(proto.num(), 1);
74   CHECK(proto.channels() == 3 || proto.channels() == 1);
75   CHECK_GT(proto.height(), 0);
76   CHECK_GT(proto.width(), 0);
77   Mat cv_img(proto.height(), proto.width(), CV_8UC3);
78   if (proto.channels() == 1) {
79     for (int c = 0; c < 3; ++c) {
80       for (int h = 0; h < cv_img.rows; ++h) {
81         for (int w = 0; w < cv_img.cols; ++w) {
82           cv_img.at<Vec3b>(h, w)[c] =
83               uint8_t(proto.data(h * cv_img.cols + w) * 255.);
84         }
85       }
86     }
87   } else {
88     for (int c = 0; c < 3; ++c) {
89       for (int h = 0; h < cv_img.rows; ++h) {
90         for (int w = 0; w < cv_img.cols; ++w) {
91           cv_img.at<Vec3b>(h, w)[c] =
92               uint8_t(proto.data((c * cv_img.rows + h) * cv_img.cols + w)
93                   * 255.);
94         }
95       }
96     }
97   }
98   CHECK(cv::imwrite(filename, cv_img));
99 }
101 void ReadProtoFromTextFile(const char* filename,
102     ::google::protobuf::Message* proto) {
103   int fd = open(filename, O_RDONLY);
104   FileInputStream* input = new FileInputStream(fd);
105   CHECK(google::protobuf::TextFormat::Parse(input, proto));
106   delete input;
107   close(fd);
110 void WriteProtoToTextFile(const Message& proto, const char* filename) {
111   int fd = open(filename, O_WRONLY);
112   FileOutputStream* output = new FileOutputStream(fd);
113   CHECK(google::protobuf::TextFormat::Print(proto, output));
114   delete output;
115   close(fd);
118 void ReadProtoFromBinaryFile(const char* filename, Message* proto) {
119   fstream input(filename, ios::in | ios::binary);
120   CHECK(proto->ParseFromIstream(&input));
123 void WriteProtoToBinaryFile(const Message& proto, const char* filename) {
124   fstream output(filename, ios::out | ios::trunc | ios::binary);
125   CHECK(proto.SerializeToOstream(&output));
128 }  // namespace caffe