uvos/SDImagePreprocess
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

inital commit

Browse Source
This commit is contained in:
uvos
2023-06-28 23:59:50 +02:00
commit 438c9d726c
19 changed files with 2169 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

295
main.cpp Normal file
View File

@ -0,0 +1,295 @@
#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"
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;
}
bool seamCarveResize(cv::Mat& image, const std::vector<Yolo::Detection>& detections, double targetAspectRatio = 1.0)
{
double aspectRatio = image.cols/static_cast<double>(image.rows);
bool vertical = false;
cv::Mat workImage;
if(aspectRatio > targetAspectRatio)
vertical = true;
int requiredLines = 0;
if(!vertical)
requiredLines = workImage.rows*targetAspectRatio - workImage.cols;
else
requiredLines = workImage.cols/targetAspectRatio - workImage.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);
int totalResizableSize = 0;
for(const std::pair<cv::Mat, bool>& slice : slices)
{
if(slice.second)
totalResizableSize += slice.first.cols;
}
std::vector<int> seamsForSlice(slices.size());
for(size_t i = 0; i < slices.size(); ++i)
{
seamsForSlice[i] = (static_cast<double>(slices[i].first.cols)/totalResizableSize)*requiredLines;
}
}
else
{
int totalResizableSize = 0;
std::vector<std::pair<cv::Mat, bool>> slices = cutImageIntoVertRegions(image, detections);
}
}
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, 4);
std::string label = detection.className + ' ' + std::to_string(detection.confidence).substr(0, 4);
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_DUPLEX, 3, 2, 0);
cv::Rect textBox(detection.box.x, detection.box.y - 80, 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, 3, cv::Scalar(0, 0, 0), 2, 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);
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)
{
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);
}
}
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;
cv::Rect crop = intRoi.getCropRectangle(detections, image.size());
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);
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;
}