lmdb.cpp
5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
#include <openbr/openbr_plugin.h>
#include "openbr/plugins/openbr_internal.h"
#include <openbr/core/qtutils.h>
#include <QFutureSynchronizer>
#include <QtConcurrentRun>
#include <QMutexLocker>
#include <QWaitCondition>
#include "caffe/util/db.hpp"
#include "caffe/util/io.hpp"
using namespace br;
class lmdbGallery : public Gallery
{
Q_OBJECT
TemplateList readBlock(bool *done)
{
*done = false;
if (!initialized) {
db = QSharedPointer<caffe::db::DB>(caffe::db::GetDB("lmdb"));
db->Open(file.name.toStdString(),caffe::db::READ);
cursor = QSharedPointer<caffe::db::Cursor>(db->NewCursor());
initialized = true;
}
caffe::Datum datum;
datum.ParseFromString(cursor->value());
cv::Mat img;
if (datum.encoded()) {
img = caffe::DecodeDatumToCVMatNative(datum);
}
else {
// create output image of appropriate size
img.create(datum.height(), datum.width(), CV_MAKETYPE(CV_8U, datum.channels()));
// copy matrix data from datum.
for (int h = 0; h < datum.height(); ++h) {
uchar* ptr = 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++] = (unsigned char)datum.data()[datum_index];
}
}
}
}
// We acquired the image data, now decode filename from db key
QString baseKey = cursor->key().c_str();
int idx = baseKey.indexOf("_");
if (idx != -1)
baseKey = baseKey.right(baseKey.size() - idx - 1);
TemplateList output;
output.append(Template(img));
output.last().file.name = baseKey;
output.last().file.set("Label", datum.label());
cursor->Next();
if (!cursor->valid())
*done = true;
return output;
}
bool initialized;
QSharedPointer<caffe::db::DB> db;
QSharedPointer<caffe::db::Cursor> cursor;
QFutureSynchronizer<void> aThread;
QMutex dataLock;
QWaitCondition dataWait;
bool should_end;
TemplateList data;
QHash<QString, int> observedLabels;
static void commitLoop(lmdbGallery * base)
{
QSharedPointer<caffe::db::Transaction> txn;
int total_count = 0;
// Acquire the lock
QMutexLocker lock(&base->dataLock);
while (true) {
// wait for data, or end signal
while(base->data.isEmpty() && !base->should_end)
base->dataWait.wait(&base->dataLock);
// If should_end, but there is still data, we need another commit
// round
if (base->should_end && base->data.isEmpty())
break;
txn = QSharedPointer<caffe::db::Transaction>(base->db->NewTransaction());
TemplateList working = base->data;
base->data.clear();
// no longer blocking dataLock
lock.unlock();
foreach(const Template &t, working) {
// add current image to transaction
caffe::Datum datum;
caffe::CVMatToDatum(t.m(), &datum);
QVariant base_label = t.file.value("Label");
QString label_str = base_label.toString();
if (!base->observedLabels.contains(label_str) )
base->observedLabels.insert(label_str, base->observedLabels.size());
datum.set_label(base->observedLabels[label_str]);
std::string out;
datum.SerializeToString(&out);
char key_cstr[256];
int len = snprintf(key_cstr, 256, "%08d_%s", total_count, qPrintable(t.file.name));
txn->Put(std::string(key_cstr, len), out);
total_count++;
}
txn->Commit();
lock.relock();
}
}
void write(const Template &t)
{
if (!initialized) {
db = QSharedPointer<caffe::db::DB> (caffe::db::GetDB("lmdb"));
db->Open(file.name.toStdString(), caffe::db::NEW);
observedLabels.clear();
initialized = true;
should_end = false;
// fire thread
aThread.clearFutures();
aThread.addFuture(QtConcurrent::run(lmdbGallery::commitLoop, this));
}
QMutexLocker lock(&dataLock);
data.append(t);
dataWait.wakeAll();
}
~lmdbGallery()
{
if (initialized) {
QMutexLocker lock(&dataLock);
should_end = true;
dataWait.wakeAll();
lock.unlock();
aThread.waitForFinished();
}
}
void init()
{
initialized = false;
should_end = false;
}
};
BR_REGISTER(Gallery, lmdbGallery)
#include "gallery/lmdb.moc"