Skip to content

Commit 823a6e4

Browse files
Nuzhny007Nuzhny007
committed
Darknet works with batch size
1 parent 9560067 commit 823a6e4

6 files changed

Lines changed: 92 additions & 66 deletions

File tree

‎example/VideoExample.cpp‎

Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,6 @@ void VideoExample::SyncProcess()
104104
bool manualMode = false;
105105
#endif
106106

107-
cv::Mat frame;
108-
109107
double freq = cv::getTickFrequency();
110108
int64 allTime = 0;
111109

@@ -118,17 +116,35 @@ void VideoExample::SyncProcess()
118116
return;
119117
}
120118

119+
FrameInfo frameInfo(m_batchSize);
120+
frameInfo.m_frames.resize(frameInfo.m_batchSize);
121+
frameInfo.m_frameInds.resize(frameInfo.m_batchSize);
122+
121123
int64 startLoopTime = cv::getTickCount();
122124

123125
for (;;)
124126
{
125-
capture >> frame;
126-
if (frame.empty())
127-
break;
127+
size_t i = 0;
128+
for (; i < m_batchSize; ++i)
129+
{
130+
capture >> frameInfo.m_frames[i].GetMatBGRWrite();
131+
if (frameInfo.m_frames[i].empty())
132+
break;
133+
frameInfo.m_frameInds[i] = framesCounter;
134+
135+
++framesCounter;
136+
if (m_endFrame && framesCounter > m_endFrame)
137+
{
138+
std::cout << "Process: riched last " << m_endFrame << " frame" << std::endl;
139+
break;
140+
}
141+
}
142+
if (i < m_batchSize)
143+
break;
128144

129145
if (!m_isDetectorInitialized || !m_isTrackerInitialized)
130146
{
131-
cv::UMat ufirst = frame.getUMat(cv::ACCESS_READ);
147+
cv::UMat ufirst = frameInfo.m_frames[0].GetUMatBGR();
132148
if (!m_isDetectorInitialized)
133149
{
134150
m_isDetectorInitialized = InitDetector(ufirst);
@@ -149,28 +165,6 @@ void VideoExample::SyncProcess()
149165
}
150166
}
151167

152-
FrameInfo frameInfo(m_batchSize);
153-
frameInfo.m_frames.emplace_back(frame);
154-
frameInfo.m_frameInds.emplace_back(framesCounter);
155-
++framesCounter;
156-
for (size_t i = 1; i < m_batchSize; ++i)
157-
{
158-
capture >> frame;
159-
if (frame.empty())
160-
break;
161-
frameInfo.m_frames.emplace_back(frame);
162-
frameInfo.m_frameInds.emplace_back(framesCounter);
163-
164-
++framesCounter;
165-
if (m_endFrame && framesCounter > m_endFrame)
166-
{
167-
std::cout << "Process: riched last " << m_endFrame << " frame" << std::endl;
168-
break;
169-
}
170-
}
171-
if (frameInfo.m_frames.size() < m_batchSize)
172-
break;
173-
174168
int64 t1 = cv::getTickCount();
175169

176170
regions_t regions;
@@ -181,12 +175,12 @@ void VideoExample::SyncProcess()
181175
allTime += t2 - t1;
182176
int currTime = cvRound(1000 * (t2 - t1) / freq);
183177

184-
for (size_t i = 0; i < m_batchSize; ++i)
178+
for (i = 0; i < m_batchSize; ++i)
185179
{
186180
DrawData(frameInfo.m_frames[i].GetMatBGR(), frameInfo.m_tracks[i], frameInfo.m_frameInds[i], currTime);
187181

188182
#ifndef SILENT_WORK
189-
cv::imshow("Video", frame);
183+
cv::imshow("Video", frameInfo.m_frames[i].GetMatBGR());
190184

191185
int waitTime = manualMode ? 0 : 1;// std::max<int>(1, cvRound(1000 / m_fps - currTime));
192186
int k = cv::waitKey(waitTime);
@@ -228,14 +222,13 @@ void VideoExample::AsyncProcess()
228222

229223
double freq = cv::getTickFrequency();
230224

231-
int framesCounter = m_startFrame + 1;
232-
233225
int64 allTime = 0;
234226
int64 startLoopTime = cv::getTickCount();
235227
size_t processCounter = 0;
236228
for (; !stopCapture.load(); )
237229
{
238230
FrameInfo& frameInfo = m_frameInfo[processCounter % 2];
231+
//std::cout << "tracking from " << (processCounter % 2) << " ind = " << processCounter << std::endl;
239232
{
240233
std::unique_lock<std::mutex> lock(frameInfo.m_mutex);
241234
if (!frameInfo.m_cond.wait_for(lock, std::chrono::milliseconds(m_captureTimeOut), [&frameInfo]{ return frameInfo.m_captured; }))
@@ -244,6 +237,7 @@ void VideoExample::AsyncProcess()
244237
break;
245238
}
246239
}
240+
//std::cout << "tracking from " << (processCounter % 2) << " in process..." << std::endl;
247241

248242
if (!m_isTrackerInitialized)
249243
{
@@ -270,12 +264,10 @@ void VideoExample::AsyncProcess()
270264
int key = 0;
271265
for (size_t i = 0; i < m_batchSize; ++i)
272266
{
273-
DrawData(frameInfo.m_frames[i].GetMatBGR(), frameInfo.m_tracks[i], framesCounter, currTime);
267+
DrawData(frameInfo.m_frames[i].GetMatBGR(), frameInfo.m_tracks[i], frameInfo.m_frameInds[i], currTime);
274268

275269
WriteFrame(writer, frameInfo.m_frames[i].GetMatBGR());
276270

277-
++framesCounter;
278-
279271
#ifndef SILENT_WORK
280272
cv::imshow("Video", frameInfo.m_frames[0].GetMatBGR());
281273

@@ -292,19 +284,16 @@ void VideoExample::AsyncProcess()
292284

293285
{
294286
std::unique_lock<std::mutex> lock(frameInfo.m_mutex);
287+
//std::cout << "tracking m_captured " << (processCounter % 2) << " - " << frameInfo.m_captured << std::endl;
288+
assert(frameInfo.m_captured);
295289
frameInfo.m_captured = false;
296290
}
297291
frameInfo.m_cond.notify_one();
298292

299293
if (key == 27)
300294
break;
301295

302-
if (m_endFrame && framesCounter > m_endFrame)
303-
{
304-
std::cout << "Process: riched last " << m_endFrame << " frame" << std::endl;
305-
break;
306-
}
307-
++processCounter;
296+
++processCounter;
308297
}
309298
stopCapture = true;
310299

@@ -342,7 +331,7 @@ void VideoExample::CaptureAndDetect(VideoExample* thisPtr, std::atomic<bool>& st
342331
for (; !stopCapture.load();)
343332
{
344333
FrameInfo& frameInfo = thisPtr->m_frameInfo[processCounter % 2];
345-
334+
//std::cout << "captured to " << (processCounter % 2) << " ind = " << processCounter << std::endl;
346335
{
347336
std::unique_lock<std::mutex> lock(frameInfo.m_mutex);
348337
if (!frameInfo.m_cond.wait_for(lock, std::chrono::milliseconds(trackingTimeOut), [&frameInfo]{ return !frameInfo.m_captured; }))
@@ -352,14 +341,16 @@ void VideoExample::CaptureAndDetect(VideoExample* thisPtr, std::atomic<bool>& st
352341
break;
353342
}
354343
}
344+
//std::cout << "capture from " << (processCounter % 2) << " in process..." << std::endl;
355345

356346
if (frameInfo.m_frames.size() < frameInfo.m_batchSize)
357347
{
358348
frameInfo.m_frames.resize(frameInfo.m_batchSize);
359349
frameInfo.m_frameInds.resize(frameInfo.m_batchSize);
360350
}
361351

362-
for (size_t i = 1; i < frameInfo.m_batchSize; ++i)
352+
size_t i = 0;
353+
for (; i < frameInfo.m_batchSize; ++i)
363354
{
364355
capture >> frameInfo.m_frames[i].GetMatBGRWrite();
365356
if (frameInfo.m_frames[i].empty())
@@ -370,8 +361,14 @@ void VideoExample::CaptureAndDetect(VideoExample* thisPtr, std::atomic<bool>& st
370361
}
371362
frameInfo.m_frameInds[i] = framesCounter;
372363
++framesCounter;
364+
365+
if (thisPtr->m_endFrame && framesCounter > thisPtr->m_endFrame)
366+
{
367+
std::cout << "Process: riched last " << thisPtr->m_endFrame << " frame" << std::endl;
368+
break;
369+
}
373370
}
374-
if (frameInfo.m_frames.size() < frameInfo.m_batchSize)
371+
if (i < frameInfo.m_batchSize)
375372
break;
376373

377374
if (!thisPtr->m_isDetectorInitialized)
@@ -392,6 +389,8 @@ void VideoExample::CaptureAndDetect(VideoExample* thisPtr, std::atomic<bool>& st
392389

393390
{
394391
std::unique_lock<std::mutex> lock(frameInfo.m_mutex);
392+
//std::cout << "capture m_captured " << (processCounter % 2) << " - " << frameInfo.m_captured << std::endl;
393+
assert(!frameInfo.m_captured);
395394
frameInfo.m_captured = true;
396395
}
397396
frameInfo.m_cond.notify_one();
@@ -438,7 +437,7 @@ void VideoExample::Tracking(FrameInfo& frame)
438437
{
439438
assert(frame.m_regions.size() == frame.m_frames.size());
440439

441-
frame.CleanåTracks();
440+
frame.CleanTracks();
442441
for (size_t i = 0; i < frame.m_frames.size(); ++i)
443442
{
444443
if (m_tracker->CanColorFrameToTrack())

‎example/VideoExample.h‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -183,30 +183,47 @@ class Frame
183183
///
184184
cv::Mat& GetMatBGRWrite()
185185
{
186+
m_umBGRGenerated = false;
187+
m_mGrayGenerated = false;
188+
m_umGrayGenerated = false;
186189
return m_mBGR;
187190
}
188191
///
189192
const cv::Mat& GetMatGray()
190193
{
191-
if (m_mGray.empty())
192-
cv::cvtColor(m_mBGR, m_mGray, cv::COLOR_BGR2GRAY);
194+
if (m_mGray.empty() || !m_mGrayGenerated)
195+
{
196+
if (m_umGray.empty() || !m_umGrayGenerated)
197+
cv::cvtColor(m_mBGR, m_mGray, cv::COLOR_BGR2GRAY);
198+
else
199+
m_mGray = m_umGray.getMat(cv::ACCESS_READ);
200+
m_mGrayGenerated = true;
201+
}
193202
return m_mGray;
194203
}
195204
///
196205
const cv::UMat& GetUMatBGR()
197206
{
198-
if (m_umBGR.empty())
207+
std::thread::id lastThreadID = std::this_thread::get_id();
208+
209+
if (m_umBGR.empty() || !m_umBGRGenerated || lastThreadID != m_umBGRThreadID)
210+
{
199211
m_umBGR = m_mBGR.getUMat(cv::ACCESS_READ);
212+
m_umBGRGenerated = true;
213+
m_umBGRThreadID = lastThreadID;
214+
}
200215
return m_umBGR;
201216
}
202217
///
203218
const cv::UMat& GetUMatGray()
204219
{
205-
if (m_umGray.empty())
220+
std::thread::id lastThreadID = std::this_thread::get_id();
221+
222+
if (m_umGray.empty() || !m_umGrayGenerated || lastThreadID != m_umGrayThreadID)
206223
{
207-
if (m_mGray.empty())
224+
if (m_mGray.empty() || !m_mGrayGenerated)
208225
{
209-
if (m_umBGR.empty())
226+
if (m_umBGR.empty() || !m_umBGRGenerated || lastThreadID != m_umGrayThreadID)
210227
cv::cvtColor(m_mBGR, m_umGray, cv::COLOR_BGR2GRAY);
211228
else
212229
cv::cvtColor(m_umBGR, m_umGray, cv::COLOR_BGR2GRAY);
@@ -215,6 +232,8 @@ class Frame
215232
{
216233
m_umGray = m_mGray.getUMat(cv::ACCESS_READ);
217234
}
235+
m_umGrayGenerated = true;
236+
m_umGrayThreadID = lastThreadID;
218237
}
219238
return m_umGray;
220239
}
@@ -224,6 +243,11 @@ class Frame
224243
cv::Mat m_mGray;
225244
cv::UMat m_umBGR;
226245
cv::UMat m_umGray;
246+
bool m_umBGRGenerated = false;
247+
bool m_mGrayGenerated = false;
248+
bool m_umGrayGenerated = false;
249+
std::thread::id m_umBGRThreadID;
250+
std::thread::id m_umGrayThreadID;
227251
};
228252

229253
///
@@ -268,7 +292,7 @@ struct FrameInfo
268292
}
269293

270294
///
271-
void CleanåTracks()
295+
void CleanTracks()
272296
{
273297
if (m_tracks.size() != m_batchSize)
274298
m_tracks.resize(m_batchSize);

‎example/examples.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,12 +628,14 @@ class YoloDarknetExample : public VideoExample
628628
config.emplace("modelConfiguration", pathToModel + "yolov4-tiny.cfg");
629629
config.emplace("modelBinary", pathToModel + "yolov4-tiny.weights");
630630
config.emplace("confidenceThreshold", "0.5");
631+
config.emplace("maxBatch", "4");
631632
break;
632633

633634
case YOLOModels::ScaledYOLOv4:
634635
config.emplace("modelConfiguration", pathToModel + "yolov4-csp.cfg");
635636
config.emplace("modelBinary", pathToModel + "yolov4-csp.weights");
636637
config.emplace("confidenceThreshold", "0.5");
638+
config.emplace("maxBatch", "2");
637639
break;
638640
}
639641
config.emplace("classNames", pathToModel + "coco.names");

‎src/Detector/YoloDarknetDetector.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,11 @@ bool YoloDarknetDetector::Init(const config_t& config)
3535
if (gpuId != config.end())
3636
currGPUID = std::max(0, std::stoi(gpuId->second));
3737

38-
m_detector = std::make_unique<Detector>(modelConfiguration->second, modelBinary->second, currGPUID);
38+
auto maxBatch = config.find("maxBatch");
39+
if (maxBatch != config.end())
40+
m_batchSize = std::max(1, std::stoi(maxBatch->second));
41+
42+
m_detector = std::make_unique<Detector>(modelConfiguration->second, modelBinary->second, currGPUID, static_cast<int>(m_batchSize));
3943
m_detector->nms = 0.2f;
4044

4145
auto classNames = config.find("classNames");
@@ -67,10 +71,6 @@ bool YoloDarknetDetector::Init(const config_t& config)
6771
if (maxCropRatio != config.end())
6872
m_maxCropRatio = std::stof(maxCropRatio->second);
6973

70-
auto maxBatch = config.find("maxBatch");
71-
if (maxBatch != config.end())
72-
m_batchSize = std::max(1, std::stoi(maxBatch->second));
73-
7474
m_classesWhiteList.clear();
7575
auto whiteRange = config.equal_range("white_list");
7676
for (auto it = whiteRange.first; it != whiteRange.second; ++it)

‎src/Detector/darknet/include/yolo_v2_class.hpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ struct bbox_t_container {
5656
#include <opencv2/imgproc/imgproc_c.h> // C
5757
#endif
5858

59-
extern "C" LIB_API int init(const char *configurationFilename, const char *weightsFilename, int gpu);
59+
extern "C" LIB_API int init(const char *configurationFilename, const char *weightsFilename, int gpu, int batch_size);
6060
extern "C" LIB_API int detect_image(const char *filename, bbox_t_container &container);
6161
extern "C" LIB_API int detect_mat(const uint8_t* data, const size_t data_length, bbox_t_container &container);
6262
extern "C" LIB_API int dispose();
@@ -72,11 +72,11 @@ class Detector {
7272
std::deque<std::vector<bbox_t>> prev_bbox_vec_deque;
7373
std::string _cfg_filename, _weight_filename;
7474
public:
75-
const int cur_gpu_id;
76-
float nms = .4;
75+
const int cur_gpu_id = 0;
76+
float nms = .4f;
7777
bool wait_stream;
7878

79-
LIB_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id = 0);
79+
LIB_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id, int batch_size);
8080
LIB_API ~Detector();
8181

8282
LIB_API std::vector<bbox_t> detect(std::string image_filename, float thresh = 0.2, bool use_mean = false);

‎src/Detector/darknet/src/yolo_v2_class.cpp‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ extern "C" {
2727
//static Detector* detector = NULL;
2828
static std::unique_ptr<Detector> detector;
2929

30-
int init(const char *configurationFilename, const char *weightsFilename, int gpu)
30+
int init(const char *configurationFilename, const char *weightsFilename, int gpu, int batch_size)
3131
{
32-
detector.reset(new Detector(configurationFilename, weightsFilename, gpu));
32+
detector.reset(new Detector(configurationFilename, weightsFilename, gpu, batch_size));
3333
return 1;
3434
}
3535

@@ -127,7 +127,8 @@ struct detector_gpu_t {
127127
unsigned int *track_id;
128128
};
129129

130-
LIB_API Detector::Detector(std::string cfg_filename, std::string weight_filename, int gpu_id) : cur_gpu_id(gpu_id)
130+
LIB_API Detector::Detector(std::string cfg_filename, std::string weight_filename, int gpu_id, int batch_size)
131+
: cur_gpu_id(gpu_id)
131132
{
132133
wait_stream = 0;
133134
#ifdef GPU
@@ -153,11 +154,11 @@ LIB_API Detector::Detector(std::string cfg_filename, std::string weight_filename
153154
char *cfgfile = const_cast<char *>(_cfg_filename.c_str());
154155
char *weightfile = const_cast<char *>(_weight_filename.c_str());
155156

156-
net = parse_network_cfg_custom(cfgfile, 1, 1);
157+
net = parse_network_cfg_custom(cfgfile, batch_size, batch_size);
157158
if (weightfile) {
158159
load_weights(&net, weightfile);
159160
}
160-
set_batch_network(&net, 1);
161+
set_batch_network(&net, batch_size);
161162
net.gpu_index = cur_gpu_id;
162163
fuse_conv_batchnorm(net);
163164

0 commit comments

Comments
 (0)