re-implement with C++ & new Smoother (38d2005a) · Commits · Ziqin Wang / ArmorDetection

.gitignore

+5 −4

Original line number	Diff line number	Diff line
		@@ -2,3 +2,4 @@
		.DS_Store
		.idea
		*.pyc
		cmake-build-debug/

CMakeLists.txt

0 → 100644

+23 −0

Original line number	Diff line number	Diff line
		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

src/ArmorDetectionApp.cpp

0 → 100644

+94 −0

Original line number	Diff line number	Diff line
		#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

src/ArmorDetectionApp.h

0 → 100644

+62 −0

Original line number	Diff line number	Diff line
		#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

src/Detector.cpp

0 → 100644

+224 −0

Original line number	Diff line number	Diff line
		#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

Admin message