SDImagePreprocess/main.cpp

#include <filesystem>
#include <iostream>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>
#include <algorithm>
#include <vector>

#include "yolo.h"
#include "log.h"
#include "options.h"
#include "utils.h"
#include "intelligentroi.h"
#include "seamcarving.h"

const Yolo::Detection* pointInDetectionHoriz(int x, const std::vector<Yolo::Detection>& detections, const Yolo::Detection* ignore = nullptr)
{
	const Yolo::Detection* inDetection = nullptr;
	for(const Yolo::Detection& detection : detections)
	{
		if(!ignore || ignore != &detection)
			continue;

		if(detection.box.x <= x && detection.box.x+detection.box.width <= x)
		{
			if(!inDetection || detection.box.br().x > inDetection->box.br().x)
			inDetection = &detection;
		}
	}
	return inDetection;
}

bool findRegionEndpointHoriz(int& x, const std::vector<Yolo::Detection>& detections, int imgSizeX)
{
	const Yolo::Detection* inDetection = pointInDetectionHoriz(x, detections);

	if(!inDetection)
	{
		const Yolo::Detection* closest = nullptr;
		for(const Yolo::Detection& detection : detections)
		{
			if(detection.box.x > x)
			{
				if(closest == nullptr || detection.box.x-x > closest->box.x-x)
					closest = &detection;
			}
		}
		if(closest)
			x = closest->box.x;
		else
			x = imgSizeX;
		return false;
	}
	else
	{
		x = inDetection->box.br().x;
		const Yolo::Detection* candidateDetection = pointInDetectionHoriz(x, detections, inDetection);
		if(candidateDetection && candidateDetection->box.br().x > x)
			return findRegionEndpointHoriz(x, detections, imgSizeX);
		else
			return true;
	}
}

std::vector<std::pair<cv::Mat, bool>> cutImageIntoHorzRegions(cv::Mat& image, const std::vector<Yolo::Detection>& detections)
{
	std::vector<std::pair<cv::Mat, bool>> out;

	for(int x = 0; x < image.cols; ++x)
	{
		int start = x;
		bool frozen = findRegionEndpointHoriz(x, detections, image.cols);

		cv::Mat slice = image(cv::Rect(start, 0, x-start, image.rows));
		out.push_back({slice, frozen});
	}

	return out;
}

const Yolo::Detection* pointInDetectionVert(int y, const std::vector<Yolo::Detection>& detections, const Yolo::Detection* ignore = nullptr)
{
	const Yolo::Detection* inDetection = nullptr;
	for(const Yolo::Detection& detection : detections)
	{
		if(!ignore || ignore != &detection)
			continue;

		if(detection.box.y <= y && detection.box.y+detection.box.height <= y)
		{
			if(!inDetection || detection.box.br().y > inDetection->box.br().y)
			inDetection = &detection;
		}
	}
	return inDetection;
}

bool findRegionEndpointVert(int& y, const std::vector<Yolo::Detection>& detections, int imgSizeY)
{
	const Yolo::Detection* inDetection = pointInDetectionVert(y, detections);

	if(!inDetection)
	{
		const Yolo::Detection* closest = nullptr;
		for(const Yolo::Detection& detection : detections)
		{
			if(detection.box.y > y)
			{
				if(closest == nullptr || detection.box.y-y > closest->box.y-y)
					closest = &detection;
			}
		}
		if(closest)
			y = closest->box.y;
		else
			y = imgSizeY;
		return false;
	}
	else
	{
		y = inDetection->box.br().y;
		const Yolo::Detection* candidateDetection = pointInDetectionVert(y, detections, inDetection);
		if(candidateDetection && candidateDetection->box.br().y > y)
			return findRegionEndpointVert(y, detections, imgSizeY);
		else
			return true;
	}
}

std::vector<std::pair<cv::Mat, bool>> cutImageIntoVertRegions(cv::Mat& image, const std::vector<Yolo::Detection>& detections)
{
	std::vector<std::pair<cv::Mat, bool>> out;

	for(int y = 0; y < image.rows; ++y)
	{
		int start = y;
		bool frozen = findRegionEndpointVert(y, detections, image.rows);

		cv::Mat slice = image(cv::Rect(0, start, image.cols, y-start));
		out.push_back({slice, frozen});
	}

	return out;
}

cv::Mat assembleFromSlicesVert(const std::vector<std::pair<cv::Mat, bool>>& slices)
{
	assert(!slices.empty());

	int rows = 0;
	for(const std::pair<cv::Mat, bool>& slice : slices)
		rows += slice.first.rows;

	cv::Mat image(slices[0].first.cols, rows, slices[0].first.type());

	Log(Log::DEBUG)<<__func__<<" assembled image size "<<image.size;

	int row = 0;
	for(const std::pair<cv::Mat, bool>& slice : slices)
	{
		cv::Rect rect(0, row, slice.first.cols, slice.first.rows);
		slice.first.copyTo(image(rect));
		row += slice.first.rows;
	}

	return image;
}

cv::Mat assembleFromSlicesHoriz(const std::vector<std::pair<cv::Mat, bool>>& slices)
{
	assert(!slices.empty());

	int cols = 0;
	for(const std::pair<cv::Mat, bool>& slice : slices)
		cols += slice.first.cols;

	cv::Mat image(cols, slices[0].first.rows, slices[0].first.type());

	int col = 0;
	for(const std::pair<cv::Mat, bool>& slice : slices)
	{
		cv::Rect rect(col, 0, slice.first.cols, slice.first.rows);
		slice.first.copyTo(image(rect));
		col += slice.first.cols;
	}

	return image;
}

bool seamCarveResize(cv::Mat& image, std::vector<Yolo::Detection> detections, double targetAspectRatio = 1.0)
{
	detections.erase(std::remove_if(detections.begin(), detections.end(), [](const Yolo::Detection& detection){return detection.priority < 3;}), detections.end());

	double aspectRatio = image.cols/static_cast<double>(image.rows);

	Log(Log::DEBUG)<<"Image size "<<image.size()<<" aspect ratio "<<aspectRatio<<" target aspect ratio "<<targetAspectRatio;

	bool vertical = false;
	if(aspectRatio > targetAspectRatio)
		vertical = true;

	int requiredLines = 0;
	if(!vertical)
		requiredLines = image.rows*targetAspectRatio - image.cols;
	else
		requiredLines = image.cols/targetAspectRatio - image.rows;

	Log(Log::DEBUG)<<__func__<<' '<<requiredLines<<" lines are required in "<<(vertical ? "vertical" : "horizontal")<<" direction";

	if(!vertical)
	{
		std::vector<std::pair<cv::Mat, bool>> slices = cutImageIntoHorzRegions(image, detections);
		Log(Log::DEBUG)<<"Image has "<<slices.size()<<" slices";
		int totalResizableSize = 0;
		for(const std::pair<cv::Mat, bool>& slice : slices)
		{
			if(slice.second)
				totalResizableSize += slice.first.cols;
		}

		if(totalResizableSize < requiredLines/10)
		{
			Log(Log::WARN)<<"Unable to seam carve as there are only "<<totalResizableSize<<" unfrozen cols";
			return false;
		}

		for(size_t i = 0; i < slices.size(); ++i)
		{
			if(slices[i].second)
			{
				int seamsForSlice = (static_cast<double>(slices[i].first.cols)/totalResizableSize)*requiredLines;
				bool ret = SeamCarving::strechImage(image, seamsForSlice, true);
				if(!ret)
					return false;
			}
		}

		image = assembleFromSlicesHoriz(slices);
	}
	else
	{
		std::vector<std::pair<cv::Mat, bool>> slices = cutImageIntoVertRegions(image, detections);
		Log(Log::DEBUG)<<"Image has "<<slices.size()<<" slices";
		int totalResizableSize = 0;
		for(const std::pair<cv::Mat, bool>& slice : slices)
		{
			if(slice.second)
				totalResizableSize += slice.first.rows;
		}

		if(totalResizableSize < requiredLines/10)
		{
			Log(Log::WARN)<<"Unable to seam carve as there are only "<<totalResizableSize<<" unfrozen rows";
			return false;
		}

		for(size_t i = 0; i < slices.size(); ++i)
		{
			if(slices[i].second)
			{
				int seamsForSlice = (static_cast<double>(slices[i].first.rows)/totalResizableSize)*requiredLines;
				bool ret = SeamCarving::strechImageVert(image, seamsForSlice, true);
				if(!ret)
					return false;
			}
		}

		image = assembleFromSlicesVert(slices);
	}
	return true;
}

void drawDebugInfo(cv::Mat &image, const cv::Rect& rect, const std::vector<Yolo::Detection>& detections)
{
	for(const Yolo::Detection& detection : detections)
	{
		cv::rectangle(image, detection.box, detection.color, 3);
		std::string label = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
		cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX, 1, 1, 0);
		cv::Rect textBox(detection.box.x, detection.box.y - 40, labelSize.width + 10, labelSize.height + 20);
		cv::rectangle(image, textBox, detection.color, cv::FILLED);
		cv::putText(image, label, cv::Point(detection.box.x + 5, detection.box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 1, 0);
	}

	cv::rectangle(image, rect, cv::Scalar(0, 0, 255), 8);
}

int main(int argc, char* argv[])
{
	Log::level = Log::INFO;

	Config config;
	argp_parse(&argp, argc, argv, 0, 0, &config);

	if(config.outputDir.empty())
	{
		Log(Log::ERROR)<<"a output path \"-o\" is required";
		return 1;
	}

	if(config.imagePaths.empty())
	{
		Log(Log::ERROR)<<"at least one input image or directory is required";
		return 1;
	}

	std::vector<std::filesystem::path> imagePaths;

	for(const std::filesystem::path& path : config.imagePaths)
		getImageFiles(path, imagePaths);

	Log(Log::DEBUG)<<"Images:";
	for(const::std::filesystem::path& path: imagePaths)
		Log(Log::DEBUG)<<path;

	if(imagePaths.empty())
	{
		Log(Log::ERROR)<<"no image was found\n";
		return 1;
	}

	Yolo yolo(config.modelPath, {640, 480}, config.classesPath, false);
	InteligentRoi intRoi(yolo);

	if(!std::filesystem::exists(config.outputDir))
	{
		if(!std::filesystem::create_directory(config.outputDir))
		{
			Log(Log::ERROR)<<"could not create directory at "<<config.outputDir;
			return 1;
		}
	}

	std::filesystem::path debugOutputPath(config.outputDir/"debug");
	if(config.debug)
	{
		if(!std::filesystem::exists(debugOutputPath))
			std::filesystem::create_directory(debugOutputPath);
	}

	for(const std::filesystem::path& path : imagePaths)
	{
		cv::Mat image = cv::imread(path);
		if(!image.data)
		{
			Log(Log::WARN)<<"could not load image "<<path<<" skipping";
			continue;
		}

		if(std::max(image.cols, image.rows) > 1024)
		{
			Log(Log::DEBUG, false)<<"Image is "<<image.size();
			if(image.cols > image.rows)
			{
				double ratio = 1024.0/image.cols;
				cv::resize(image, image, {1024, static_cast<int>(image.rows*ratio)}, 0, 0, cv::INTER_CUBIC);
			}
			else
			{
				double ratio = 1024.0/image.rows;
				cv::resize(image, image, {static_cast<int>(image.cols*ratio), 1024}, 0, 0, cv::INTER_CUBIC);
			}
			Log(Log::DEBUG)<<" resized to "<<image.size();
		}

		std::vector<Yolo::Detection> detections = yolo.runInference(image);

		Log(Log::DEBUG)<<"Got "<<detections.size()<<" detections for "<<path;
		for(const Yolo::Detection& detection : detections)
			Log(Log::DEBUG)<<detection.class_id<<": "<<detection.className<<" at "<<detection.box<<" with prio "<<detection.priority;

		if(config.seamCarving)
		{
			Log(Log::INFO)<<"Trying seam resize";
			seamCarveResize(image, detections, 1);
		}

		cv::Rect crop = intRoi.getCropRectangle(detections, image.size());

		if(config.debug)
		{
			cv::Mat debugImage = image.clone();
			drawDebugInfo(debugImage, crop, detections);
			bool ret = cv::imwrite(debugOutputPath/path.filename(), debugImage);
			if(!ret)
				Log(Log::WARN)<<"could not save debug image to "<<debugOutputPath/path.filename()<<" skipping";
		}

		cv::Mat croppedImage = image(crop);
		cv::Mat resizedImage;
		cv::resize(croppedImage, resizedImage, {512, 512}, 0, 0, cv::INTER_CUBIC);
		bool ret = cv::imwrite(config.outputDir/path.filename(), resizedImage);
		if(!ret)
			Log(Log::WARN)<<"could not save image to "<<config.outputDir/path.filename()<<" skipping";
	}
	return 0;
}