tgaの読み書きに対応

.gitmodules

@@ -6,3 +6,6 @@
 	url = https://github.com/lltcggie/caffe.git
 	branch = support_windows
 	ignore = dirty
+[submodule "stb"]
+	path = stb
+	url = https://github.com/nothings/stb

common/waifu2x.cpp

@@ -10,10 +10,13 @@
 #include <chrono>
 #include <cuda_runtime.h>
 
+#define STB_IMAGE_IMPLEMENTATION
+#include <stb_image.h>
+#define STB_IMAGE_WRITE_IMPLEMENTATION
+#include <stb_image_write.h>
+
 #if defined(WIN32) || defined(WIN64)
 #include <Windows.h>
-
-#undef LoadImage
 #endif
 
 #ifdef _MSC_VER
@@ -66,6 +69,28 @@ static std::once_flag waifu2x_cuda_once_flag;
 		} \
 	} while (0)
 
+namespace
+{
+	class IgnoreErrorCV
+	{
+	private:
+		static int handleError(int status, const char* func_name,
+			const char* err_msg, const char* file_name,
+			int line, void* userdata)
+		{
+			return 0;
+		}
+
+	public:
+		IgnoreErrorCV()
+		{
+			cv::redirectError(handleError);
+		}
+	};
+
+	IgnoreErrorCV g_IgnoreErrorCV;
+}
+
 Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr)
 {
 }
@@ -153,12 +178,24 @@ Waifu2x::eWaifu2xCudaError Waifu2x::can_use_CUDA()
 	return CudaFlag;
 }
 
+cv::Mat Waifu2x::LoadMat(const std::string &path)
+{
+	cv::Mat mat;
+	LoadMat(mat, path);
+
+	return mat;
+}
+
 // 画像を読み込んで値を0.0f～1.0fの範囲に変換
-Waifu2x::eWaifu2xError Waifu2x::LoadImage(cv::Mat &float_image, const std::string &input_file)
+Waifu2x::eWaifu2xError Waifu2x::LoadMat(cv::Mat &float_image, const std::string &input_file)
 {
 	cv::Mat original_image = cv::imread(input_file, cv::IMREAD_UNCHANGED);
 	if (original_image.empty())
-		return eWaifu2xError_FailedOpenInputFile;
+	{
+		const eWaifu2xError ret = LoadMatBySTBI(original_image, input_file);
+		if (ret != eWaifu2xError_OK)
+			return ret;
+	}
 
 	cv::Mat convert;
 	original_image.convertTo(convert, CV_32F, 1.0 / 255.0);
@@ -188,6 +225,62 @@ Waifu2x::eWaifu2xError Waifu2x::LoadImage(cv::Mat &float_image, const std::strin
 	return eWaifu2xError_OK;
 }
 
+Waifu2x::eWaifu2xError Waifu2x::LoadMatBySTBI(cv::Mat &float_image, const std::string &input_file)
+{
+	int x, y, comp;
+	stbi_uc *data = stbi_load(input_file.c_str(), &x, &y, &comp, 4);
+	if (!data)
+		return eWaifu2xError_FailedOpenInputFile;
+
+	int type = 0;
+	switch (comp)
+	{
+	case 1:
+	case 3:
+	case 4:
+		type = CV_MAKETYPE(CV_8U, comp);
+		break;
+
+	default:
+		return eWaifu2xError_FailedOpenInputFile;
+	}
+
+	float_image = cv::Mat(cv::Size(x, y), type);
+
+	const auto LinePixel = float_image.step1() / float_image.channels();
+	const auto Channel = float_image.channels();
+	const auto Width = float_image.size().width;
+	const auto Height = float_image.size().height;
+
+	assert(x == Width);
+	assert(y == Height);
+	assert(Channel == comp);
+
+	auto ptr = float_image.data;
+	for (int i = 0; i < y; i++)
+	{
+		for (int j = 0; j < x; j++)
+		{
+			for (int ch = 0; ch < Channel; ch++)
+				ptr[(i * LinePixel + j) * comp + ch] = data[(i * x + j) * comp + ch];
+		}
+	}
+
+	stbi_image_free(data);
+
+	if (comp >= 3)
+	{
+		// RGBだからBGRに変換
+		for (int i = 0; i < y; i++)
+		{
+			for (int j = 0; j < x; j++)
+				std::swap(ptr[(i * LinePixel + j) * comp + 0], ptr[(i * LinePixel + j) * comp + 2]);
+		}
+	}
+
+	return eWaifu2xError_OK;
+}
+
 // 画像から輝度の画像を取り出す
 Waifu2x::eWaifu2xError Waifu2x::CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im)
 {
@@ -874,6 +967,68 @@ void Waifu2x::destroy()
 	is_inited = false;
 }
 
+Waifu2x::eWaifu2xError Waifu2x::WriteMat(const cv::Mat &im, const std::string &output_file)
+{
+	const boost::filesystem::path ip(output_file);
+	const std::string ext = ip.extension().string();
+
+	const bool isJpeg = boost::iequals(ext, ".jpg") || boost::iequals(ext, ".jpeg");
+
+	if (boost::iequals(ext, ".tga"))
+	{
+		unsigned char *data = im.data;
+
+		std::vector<unsigned char> rgbimg;
+		if (im.channels() >= 3 || im.step1() != im.size().width * im.channels()) // RGB用バッファにコピー(あるいはパディングをとる)
+		{
+			const auto Line = im.step1();
+			const auto Channel = im.channels();
+			const auto Width = im.size().width;
+			const auto Height = im.size().height;
+
+			rgbimg.resize(Width * Height * Channel);
+
+			const auto Stride = Width * Channel;
+			for (int i = 0; i < Height; i++)
+				memcpy(rgbimg.data() + Stride * i, im.data + Line * i, Stride);
+
+			data = rgbimg.data();
+		}
+
+		if (im.channels() >= 3) // BGRをRGBに並び替え
+		{
+			const auto Line = im.step1();
+			const auto Channel = im.channels();
+			const auto Width = im.size().width;
+			const auto Height = im.size().height;
+
+			auto ptr = rgbimg.data();
+			for (int i = 0; i < Height; i++)
+			{
+				for (int j = 0; j < Width; j++)
+					std::swap(ptr[(i * Width + j) * Channel + 0], ptr[(i * Width + j) * Channel + 2]);
+			}
+		}
+
+		if(!stbi_write_tga(output_file.c_str(), im.size().width, im.size().height, im.channels(), data))
+			return eWaifu2xError_FailedOpenOutputFile;
+
+		return eWaifu2xError_OK;
+	}
+
+	try
+	{
+		if (cv::imwrite(output_file, im))
+			return eWaifu2xError_OK;
+
+	}
+	catch (...)
+	{
+	}
+
+	return eWaifu2xError_FailedOpenOutputFile;
+}
+
 Waifu2x::eWaifu2xError Waifu2x::waifu2x(const std::string &input_file, const std::string &output_file,
 	const waifu2xCancelFunc cancel_func)
 {
@@ -883,7 +1038,7 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const std::string &input_file, const std
 		return eWaifu2xError_NotInitialized;
 
 	cv::Mat float_image;
-	ret = LoadImage(float_image, input_file);
+	ret = LoadMat(float_image, input_file);
 	if (ret != eWaifu2xError_OK)
 		return ret;
 
@@ -1018,8 +1173,9 @@ Waifu2x::eWaifu2xError Waifu2x::waifu2x(const std::string &input_file, const std
 	process_image.convertTo(write_iamge, CV_8U, 255.0);
 	process_image.release();
 
-	if (!cv::imwrite(output_file, write_iamge))
-		return eWaifu2xError_FailedOpenOutputFile;
+	ret = WriteMat(write_iamge, output_file);
+	if (ret != eWaifu2xError_OK)
+		return ret;
 
 	write_iamge.release();
 

common/waifu2x.h

@@ -91,7 +91,8 @@ class Waifu2x
 	float *output_block;
 
 private:
-	eWaifu2xError LoadImage(cv::Mat &float_image, const std::string &input_file);
+	static eWaifu2xError LoadMat(cv::Mat &float_image, const std::string &input_file);
+	static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::string &input_file);
 	eWaifu2xError CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im);
 	eWaifu2xError PaddingImage(const cv::Mat &input, cv::Mat &output);
 	eWaifu2xError Zoom2xAndPaddingImage(const cv::Mat &input, cv::Mat &output, cv::Size_<int> &zoom_size);
@@ -100,6 +101,7 @@ class Waifu2x
 	eWaifu2xError LoadParameterFromJson(boost::shared_ptr<caffe::Net<float>> &net, const std::string &model_path, const std::string &param_path);
 	eWaifu2xError SetParameter(caffe::NetParameter &param) const;
 	eWaifu2xError ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im);
+	eWaifu2xError WriteMat(const cv::Mat &im, const std::string &output_file);
 
 public:
 	Waifu2x();
@@ -119,4 +121,6 @@ class Waifu2x
 		const waifu2xCancelFunc cancel_func = nullptr);
 
 	const std::string& used_process() const;
+
+	static cv::Mat LoadMat(const std::string &path);
 };

waifu2x-caffe-gui/Source.cpp

@@ -264,7 +264,7 @@ class DialogEvent
 			{
 				if (!boost::filesystem::is_directory(p) && std::find(extList.begin(), extList.end(), p.extension().string()) != extList.end())
 				{
-					auto mat = cv::imread(p.string(), cv::IMREAD_UNCHANGED);
+					auto mat = Waifu2x::LoadMat(p.string());
 					if (mat.empty())
 						continue;
 
@@ -277,7 +277,7 @@ class DialogEvent
 		}
 		else
 		{
-			auto mat = cv::imread(input_path.string(), cv::IMREAD_UNCHANGED);
+			auto mat = Waifu2x::LoadMat(input_path.string());
 			if (mat.empty())
 				return;
 

waifu2x-caffe-gui/waifu2x-caffe-gui.vcxproj

@@ -41,12 +41,12 @@
   <PropertyGroup Label="UserMacros" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
     <LinkIncremental>true</LinkIncremental>
-    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)include;$(IncludePath)</IncludePath>
+    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath)</IncludePath>
     <LibraryPath>$(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath)</LibraryPath>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
     <LinkIncremental>false</LinkIncremental>
-    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)include;$(IncludePath)</IncludePath>
+    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath)</IncludePath>
     <LibraryPath>$(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath)</LibraryPath>
   </PropertyGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">

waifu2x-caffe/waifu2x-caffe.vcxproj

@@ -41,12 +41,12 @@
   <PropertyGroup Label="UserMacros" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
     <LinkIncremental>true</LinkIncremental>
-    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)include;$(IncludePath)</IncludePath>
+    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath)</IncludePath>
     <LibraryPath>$(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath)</LibraryPath>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
     <LinkIncremental>false</LinkIncremental>
-    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)include;$(IncludePath)</IncludePath>
+    <IncludePath>$(SolutionDir)caffe\build\include;$(SolutionDir)caffe\3rdparty\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v7.0\include;C:\boost_1_56_0;C:\opencv249\build\include;$(SolutionDir)rapidjson\include;$(SolutionDir)stb;$(SolutionDir)include;$(IncludePath)</IncludePath>
     <LibraryPath>$(SolutionDir)caffe\build\lib;$(SolutionDir)caffe\3rdparty\lib;C:\boost_1_56_0\lib64-msvc-12.0;$(LibraryPath)</LibraryPath>
   </PropertyGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">