Commit 38d2005a authored by Jeeken's avatar Jeeken
Browse files

re-implement with C++ & new Smoother

parent 600a3538
.o
.DS_Store
.idea
*.pyc
.o
.DS_Store
.idea
*.pyc
cmake-build-debug/
cmake_minimum_required(VERSION 3.1)
project(ArmorDetection)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_FLAGS_RELEASE "-O3")
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})
set(SOURCE_FILES
src/main.cpp
src/Detector.cpp
src/Detector.h
src/ArmorDetectionApp.cpp
src/ArmorDetectionApp.h
src/utils.cpp
src/utils.h
src/Smoother.cpp
src/Smoother.h
src/FixedQueue.h)
add_executable(${PROJECT_NAME} ${SOURCE_FILES})
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS})
\ No newline at end of file
#include "ArmorDetectionApp.h"
#include "utils.h"
using namespace ArmorDetection;
void ArmorDetectionApp::drawTarget(Mat &img, const Point &target) {
Scalar aimColor;
if (img.channels() == 1)
aimColor = 255;
else
aimColor = color == BarColor::BLUE ? MarkerBgrColor::RED : MarkerBgrColor::GREEN;
cv::circle(img, target, 20, aimColor, 2);
cv::drawMarker(img, target, aimColor, cv::MARKER_CROSS, 80, 2);
}
ImgArmorDetectionApp::ImgArmorDetectionApp(BarColor color, const Size &frameSize,
const string &folder, const string &ext, bool debug)
: ArmorDetectionApp(color, debug),
frameSize(frameSize),
files(utils::getFilesFromFolder(folder, ext)) {}
void ImgArmorDetectionApp::run() {
for (string &filename: files) {
std::cout << "Image: " << filename << std::endl;
Mat img;
cv::resize(cv::imread(filename), img, frameSize);
cv::imshow("Origin", img);
Point target;
if (detector.target(img, target)) ArmorDetectionApp::drawTarget(img, target);
if (debug)
for (auto &imgAndTitle: detector.getDebugImgs())
cv::imshow(std::get<0>(imgAndTitle), std::get<1>(imgAndTitle));
cv::imshow("Aimed", img);
if (cv::waitKey() == 'q') {
cv::destroyAllWindows();
break;
}
}
}
VideoArmorDetectionApp::VideoArmorDetectionApp(BarColor color, const Size &frameSize, const string &file, bool debug)
: ArmorDetectionApp(color, debug),
frameSize(frameSize),
file(file),
smoother(frameSize) {
if (!utils::fileExists(file))
throw std::invalid_argument("The file does not exists!");
}
void VideoArmorDetectionApp::run() {
auto cap = cv::VideoCapture(file);
while (cap.isOpened()) {
Mat frame;
if (!cap.read(frame)) break;
cv::resize(frame, frame, frameSize);
cv::imshow("Original", frame);
Point target;
bool hasTarget = detector.target(frame, target);
if (debug)
for (auto &debugInfo: detector.getDebugImgs())
cv::imshow(std::get<0>(debugInfo), std::get<1>(debugInfo));
if (hasTarget) {
std::cout << "Target: " << target << std::endl;
cv::circle(frame, target, 4, MarkerBgrColor::GREEN, -1);
} else {
std::cout << "No Target" << std::endl;
}
bool hasValidTarget = smoother.smooth(hasTarget, target);
if (hasValidTarget) {
ArmorDetectionApp::drawTarget(frame, target);
std::cout << "Smoothed: " << target << std::endl;
} else {
std::cout << "No Smoothed" << std::endl;
}
std::cout << "--------------------" << std::endl;
cv::imshow("Aimed", frame);
if (cv::waitKey(1) == 'q') break;
}
}
\ No newline at end of file
#ifndef ARMORDETECTION_APP_H
#define ARMORDETECTION_APP_H
#include "Detector.h"
#include "Smoother.h"
namespace ArmorDetection {
class ArmorDetectionApp {
public:
explicit ArmorDetectionApp(BarColor color, bool debug = false)
: detector(color, debug),
color(color),
debug(debug) {}
virtual ~ArmorDetectionApp() {
cv::destroyAllWindows();
}
virtual void run() = 0;
protected:
Detector detector;
BarColor color;
bool debug;
void drawTarget(Mat &img, const Point &target);
};
class ImgArmorDetectionApp final : public ArmorDetectionApp {
public:
ImgArmorDetectionApp(BarColor color, const Size &frameSize, const string &folder = ".",
const string &ext = "jpg", bool debug = false);
void run() override;
private:
Size frameSize;
std::forward_list<string> files;
};
class VideoArmorDetectionApp final : public ArmorDetectionApp {
public:
VideoArmorDetectionApp(BarColor color, const Size &frameSize, const string &file, bool debug = false);
void run() override;
private:
Size frameSize;
string file;
Smoother smoother;
};
} // namespace ArmorDetection
using ArmorDetection::ImgArmorDetectionApp;
using ArmorDetection::VideoArmorDetectionApp;
using ArmorDetection::BarColor;
#endif //ARMORDETECTION_APP_H
#include "Detector.h"
#include "utils.h"
using namespace ArmorDetection;
bool Detector::target(const Mat &frame, Point &targetCenter) {
debugImgs.clear();
currentFrame = &frame;
// FIXME
Mat tmp;
cv::pyrDown(frame, tmp);
cv::pyrUp(tmp, frame);
Mat light, halo;
hsv(frame, light, halo);
Point haloCenter;
float haloRadius;
std::tie(haloCenter, haloRadius) = haloCircle(halo);
auto lightsInfo = getLights(light);
// debug
addDebugImg("Halo", [&halo, &haloCenter, haloRadius]() {
cv::circle(halo, haloCenter, utils::round(haloRadius), Scalar(255), 2);
return halo;
});
addDebugImg("Light", [&light, &haloCenter, haloRadius]() {
cv::circle(light, haloCenter, utils::round(haloRadius), Scalar(255), 2);
return light;
});
auto lightsInside = selectLightsInHalo(lightsInfo, haloCenter, haloRadius);
auto verticalLights = selectVerticalLights(lightsInside);
return getArmor(verticalLights, targetCenter);
}
void Detector::hsv(const Mat &frame, Mat &lightArea, Mat &haloArea) {
Mat hsvFrame;
cv::cvtColor(frame, hsvFrame, cv::COLOR_BGR2HSV);
if (color == BarColor::RED) {
Mat hsvSpace[3];
cv::split(hsvFrame, hsvSpace);
Mat &h = hsvSpace[0];
h += 107;
h -= (h / 180) * 180;
cv::merge(hsvSpace, 3, hsvFrame);
addDebugImg("Red2Blue", hsvFrame);
}
unsigned char lowerLight[3] = {90, 0, 200};
unsigned char upperLight[3] = {120, 180, 255};
unsigned char lowerHalo[3] = {90, 120, 185};
unsigned char upperHalo[3] = {120, 255, 255};
cv::inRange(hsvFrame, Mat(1, 3, CV_8UC1, lowerLight), Mat(1, 3, CV_8UC1, upperLight), lightArea);
cv::inRange(hsvFrame, Mat(1, 3, CV_8UC1, lowerHalo), Mat(1, 3, CV_8UC1, upperHalo), haloArea);
}
std::forward_list<RotatedRect> Detector::getLights(Mat &binaryImg) {
std::vector<Mat> contours;
cv::findContours(binaryImg, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
float verySmallArea = binaryImg.size().width / 30.0f;
auto notVerySmall = [verySmallArea](RotatedRect rect) {
return rect.size.area() > verySmallArea;
};
std::function<RotatedRect(Mat)> minAreaRect = cv::minAreaRect; // for type matching
return utils::filter(utils::map(contours, minAreaRect), notVerySmall);
}
std::tuple<Point, float> Detector::haloCircle(Mat &binaryImg) {
int width = binaryImg.size().width, height = binaryImg.size().height;
std::vector<Mat> contours;
cv::findContours(binaryImg, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
float verySmallArea = width / 30.0f;
auto selectedContours = utils::filter(contours, [verySmallArea](Mat contour) {
Mat xy[2];
cv::split(contour, xy);
double xMin, xMax, yMin, yMax;
cv::minMaxIdx(xy[0], &xMin, &xMax);
cv::minMaxIdx(xy[1], &yMin, &yMax);
return (xMax - xMin) * (yMax - yMin) > verySmallArea;
});
Point center;
float radius;
if (!selectedContours.empty()) {
cv::Point2f originalCenter;
Mat combinedHaloPoints;
for (const Mat &contour: selectedContours)
for (auto p = contour.begin<cv::Vec2i>(); p != contour.end<cv::Vec2i>(); ++p)
combinedHaloPoints.push_back(*p);
cv::minEnclosingCircle(combinedHaloPoints, originalCenter, radius);
center = originalCenter;
if (radius > width * 0.25f)
radius = width * 0.25f; // FIXME in python
else if (width / 15.0f < radius && radius < width / 4.0f)
radius *= 1.8f;
else
radius = width / 15.0f;
} else {
center = Point(width / 2, height / 2);
radius = width / 3.0f;
}
return std::make_tuple(center, radius);
};
std::forward_list<RotatedRect> Detector::selectLightsInHalo(
std::forward_list<RotatedRect> &lights, const Point &haloCenter, float haloRadius) {
int x = haloCenter.x, y = haloCenter.y;
auto lightsInside = utils::filter(lights, [x, y, haloRadius](const RotatedRect &light) {
return utils::sqr(light.center.x - x) + utils::sqr(light.center.y - y) < utils::sqr(haloRadius);
});
addDebugImg("Lights Inside", [this, &lightsInside]() {
Mat show = this->currentFrame->clone();
for (auto &light: lightsInside)
cv::ellipse(show, light, MarkerBgrColor::GREEN, 2);
return show;
});
return lightsInside;
}
std::forward_list<RotatedRect> Detector::selectVerticalLights(std::forward_list<RotatedRect> &lights) {
std::function<RotatedRect(RotatedRect)> normalizeAngle = [](const RotatedRect &light) {
return light.size.width > light.size.height
? RotatedRect(light.center, cv::Point2f(light.size.height, light.size.width), light.angle + 90)
: light;
};
auto angleNormalizedLights = utils::map(lights, normalizeAngle);
auto verticalLights = utils::filter(angleNormalizedLights, [](const RotatedRect &light) {
return -25 < light.angle && light.angle < 25;
});
addDebugImg("Vertical Lights", [this, &verticalLights]() {
Mat show = this->currentFrame->clone();
// FIXME in python, useless rand_color
for (auto &light: verticalLights)
cv::ellipse(show, light, MarkerBgrColor::GREEN, 2);
return show;
});
return verticalLights;
}
float parallel(RotatedRect light1, RotatedRect light2) {
return std::fabs(light1.angle - light2.angle) / 15.0f;
};
float shapeSimilarity(RotatedRect light1, RotatedRect light2) {
float w1 = light1.size.width, h1 = light1.size.height;
float w2 = light2.size.width, h2 = light2.size.height;
float minWidth = std::min(w1, w2), minHeight = std::min(h1, h2);
return std::fabs(w1 - w2) / minWidth + 0.33333f * std::fabs(h1 - h2) / minHeight; // FIXME in python
};
float squareRatio(RotatedRect light1, RotatedRect light2) {
float x1 = light1.center.x, y1 = light1.center.y;
float x2 = light2.center.x, y2 = light2.center.y;
float armorWidth = std::sqrt(utils::sqr(x1 - x2) + utils::sqr(y1 - y2));
float armorHeight = 0.5f * (light1.size.height + light2.size.height);
float ratio = armorWidth / armorHeight;
// FIXME in python: useless abs()
return ratio > 0.85 ? utils::sqr(ratio - 2.5f) : 1e6f;
};
float yDis(RotatedRect light1, RotatedRect light2) {
float y1 = light1.center.y, y2 = light2.center.y;
// FIXME in python: y coordinates may be negetive
return std::fabs((y1 - y2) / std::min(y1, y2));
};
bool Detector::getArmor(const std::forward_list<RotatedRect> &lights, Point &target) {
std::vector<RotatedRect> candidates;
for (auto &light: lights)
candidates.push_back(light);
if (candidates.size() < 2) return false;
typedef std::tuple<float, size_t, size_t> Result; // score, index1, index2
size_t iEnd = candidates.size() - 1, jEnd = candidates.size();
std::vector<Result> results(iEnd * jEnd / 2);
auto p = results.begin();
for (size_t i = 0; i < iEnd; ++i) {
for (size_t j = i + 1; j < jEnd; ++j) {
auto &l1 = candidates[i], &l2 = candidates[j];
float score = squareRatio(l1, l2) * 5.0f
+ yDis(l1, l2) * 20.0f
+ shapeSimilarity(l1, l2) * 3.0f
+ parallel(l1, l2) * 1.2f;
*(p++) = std::make_tuple(score, i, j);
}
}
float winnerScore;
size_t index1, index2;
std::tie(winnerScore, index1, index2) = *std::min_element(results.begin(), results.end(),
[](const Result &a, const Result &b) {
return std::get<0>(a) < std::get<0>(b);
});
debugPrint("score: " + std::to_string(winnerScore));
if (winnerScore > 7) return false;
cv::Point2f p1 = candidates[index1].center, p2 = candidates[index2].center;
addDebugImg("Selected Pair", [this, &p1, &p2]() {
Mat show = this->currentFrame->clone();
cv::drawMarker(show, Point(p1), MarkerBgrColor::GREEN, cv::MARKER_CROSS, 15, 2);
cv::drawMarker(show, Point(p2), MarkerBgrColor::GREEN, cv::MARKER_CROSS, 15, 2);
return show;
});
target = (p1 + p2) / 2;
return true;
}
\ No newline at end of file
#ifndef ARMORDETECTION_DETECTOR_H
#define ARMORDETECTION_DETECTOR_H
#include <forward_list>
#include <functional>
#include <iostream>
#include <tuple>
#include <opencv2/opencv.hpp>
namespace ArmorDetection {
using std::string;
using cv::Mat;
using cv::Point;
using cv::Scalar;
using cv::Size;
using cv::RotatedRect;
enum BarColor {
BLUE, RED
};
namespace MarkerBgrColor {
const Scalar BLUE(255, 0, 0);
const Scalar GREEN(0, 255, 0);
const Scalar RED(0, 0, 255);
const Scalar YELLOW(0, 255, 255);
const Scalar PURPLE(255, 0, 255);
}
class Detector final {
public:
Detector(BarColor color, bool debug) : color(color), debug(debug), currentFrame(nullptr) {}
void hsv(const Mat &frame, Mat &lightArea, Mat &rightArea);
bool target(const Mat &frame, Point &targetCenter);
std::forward_list<RotatedRect> getLights(Mat &binaryImg);
std::tuple<Point, float> haloCircle(Mat &binaryImg);
std::forward_list<RotatedRect> selectLightsInHalo(
std::forward_list<RotatedRect> &lights, const Point &haloCenter, float haloRadius);
std::forward_list<RotatedRect> selectVerticalLights(std::forward_list<RotatedRect> &lights);
bool getArmor(const std::forward_list<RotatedRect> &lights, Point &target);
const std::forward_list<std::tuple<string, Mat>> &getDebugImgs() const {
return debugImgs;
}
private:
BarColor color;
bool debug;
const Mat *currentFrame;
std::forward_list<std::tuple<string, Mat>> debugImgs;
void debugPrint(const string &msg) const {
if (debug) std::cerr << "Debug>>> " << msg << std::endl;
}
void addDebugImg(const string &title, const Mat &img) {
if (debug) debugImgs.push_front(std::make_tuple(title, img));
}
void addDebugImg(const string &title, const std::function<Mat()> &getImg) {
if (debug) debugImgs.push_front(std::make_tuple(title, getImg()));
}
};
}
#endif //ARMORDETECTION_DETECTOR_H
#ifndef ARMORDETECTION_FIXEDQUEUE_H
#define ARMORDETECTION_FIXEDQUEUE_H
#include <cassert>
#include <cstddef>
template<typename T>
class FixedQueue {
public:
explicit FixedQueue(size_t maxSize) : MAX_SIZE(maxSize) {
head = rear = size = 0;
data = new T[MAX_SIZE];
}
virtual ~FixedQueue() {
delete[] data;
}
void enQueue(const T &e) {
data[rear] = e;
if (isFull())
head = (head + 1) % MAX_SIZE;
else
++size;
rear = (rear + 1) % MAX_SIZE;
}
void deQueue() {
// assert(!isEmpty());
head = (head + 1) % MAX_SIZE;
--size;
}
bool isEmpty() const {
return size == 0;
}
bool isFull() const {
return size == MAX_SIZE;
}
protected:
T *data;
size_t head, rear, size;
const size_t MAX_SIZE;
};
#endif //ARMORDETECTION_FIXEDQUEUE_H
#include "Smoother.h"
#include "utils.h"
using namespace ArmorDetection;
using utils::sqr;
Smoother::Smoother(cv::Size shape)
: farEnough(std::sqrt(float(sqr(shape.width) + sqr(shape.height))) / 12), // diagonal / 12
recentPoints(CNT_THRESHOLD) {}
bool Smoother::smooth(bool hasValue, cv::Point &target) {
if (recentPoints.isEmpty()) {
if (hasValue) recentPoints.enQueue(target);
return hasValue; // no recent values, adopt new target
}
// recentPoints not empty
cv::Point recent = recentPoints.weightedAvg();
if (hasValue && cv::norm(recent - target) < farEnough) {
recentPoints.enQueue(target);
target = recentPoints.weightedAvg();
} else {
target = recent;
recentPoints.deQueue();
}
return true;
}
cv::Point Smoother::RecentPointsQueue::weightedAvg() const {
// assert(!isEmpty());
size_t p;
int cnt, sum;
cv::Point point(0, 0);
for (p = head, cnt = 1, sum = 0; cnt <= size; p = (p + 1) % MAX_SIZE, sum += cnt++)
point += data[p] * cnt;
return point / sum;
}
\ No newline at end of file
#ifndef ARMORDETECTION_SMOOTHER_H
#define ARMORDETECTION_SMOOTHER_H
#include <opencv2/opencv.hpp>
#include "FixedQueue.h"
namespace ArmorDetection {
class Smoother {
public:
explicit Smoother(cv::Size shape);