Mirror/waifu2x-caffe
1
0
Fork 0
mirror of https://github.com/lltcggie/waifu2x-caffe.git synced 2025-06-26 05:32:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

使用するGPUデバイスの指定に対応 #29

Browse Source
This commit is contained in:
lltcggie 2016-06-08 11:11:36 +09:00
parent 1caa75ef4e
commit f976d5ef64
9 changed files with 164 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
bin/lang/english.json
View File

@ -90,5 +90,6 @@
"MessageNoOverwrite":"Already file exists in the destination: %s",
"MessageCudaOldDeviceError":"Your CUDA device too old.\r\nPlease use Compute Capability 2.0 or more devices",
"OK":"OK",
"Cancel":"Cancel"
"Cancel":"Cancel",
"IDC_STATIC_USE_GPU_NO":"Use GPU No"
}

3
bin/lang/japanese.json
View File

@ -90,5 +90,6 @@
"MessageNoOverwrite":"出力先に既にファイルが存在します: %s",
"MessageCudaOldDeviceError":"CUDAデバイスが古いです。\r\nCompute Capability 2.0以上のデバイスを使用してください",
"OK":"OK",
"Cancel":"キャンセル"
"Cancel":"キャンセル",
"IDC_STATIC_USE_GPU_NO":"使用GPU No"
}

68
common/waifu2x.cpp
View File

@ -172,6 +172,41 @@ namespace
};
IgnoreErrorCV g_IgnoreErrorCV;
class CudaDeviceSet
{
private:
int orgDevice;
bool mIsSet;
public:
CudaDeviceSet(const std::string &process, const int devno) : orgDevice(0), mIsSet(false)
{
if (process == "gpu" || process == "cudnn")
{
int count = 0;
if (cudaGetDeviceCount(&count) != CUDA_SUCCESS)
return;
if (devno >= count || count < 0)
return;
if (cudaGetDevice(&orgDevice) != CUDA_SUCCESS)
return;
if (cudaSetDevice(devno) != CUDA_SUCCESS)
return;
mIsSet = true;
}
}
~CudaDeviceSet()
{
if (mIsSet)
cudaSetDevice(orgDevice);
}
};
}
template<typename BufType>
@ -323,7 +358,7 @@ static Waifu2x::eWaifu2xError writeProtoBinary(const ::google::protobuf::Message
}
Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr)
Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr), gpu_no(0)
{
}
@ -928,6 +963,8 @@ Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr<caffe::N
// ネットワークを使って画像を再構築する
Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im)
{
CudaDeviceSet devset(process, gpu_no);
const auto Height = im.size().height;
const auto Width = im.size().width;
const auto Line = im.step1();
@ -1134,7 +1171,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel,
const boost::optional<double> ScaleRatio, const boost::optional<int> ScaleWidth, const boost::optional<int> ScaleHeight,
const boost::filesystem::path &ModelDir, const std::string &Process,
const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize)
const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize, const int GPUNo)
{
Waifu2x::eWaifu2xError ret;
@ -1175,6 +1212,7 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
crop_size = CropSize;
batch_size = BatchSize;
gpu_no = GPUNo;
inner_padding = layer_num;
outer_padding = 1;
@ -1185,17 +1223,6 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
output_block_size = crop_size + (inner_padding + outer_padding - layer_num) * 2;
std::call_once(waifu2x_once_flag, [argc, argv]()
{
assert(argc >= 1);
int tmpargc = 1;
char* tmpargvv[] = { argv[0] };
char** tmpargv = tmpargvv;
// glog“™‚Ì<E2809A>‰Šú‰»
caffe::GlobalInit(&tmpargc, &tmpargv);
});
const auto cuDNNCheckStartTime = std::chrono::system_clock::now();
if (process == "gpu")
@ -1209,6 +1236,19 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
const auto cuDNNCheckEndTime = std::chrono::system_clock::now();
CudaDeviceSet devset(process, gpu_no);
std::call_once(waifu2x_once_flag, [argc, argv]()
{
assert(argc >= 1);
int tmpargc = 1;
char* tmpargvv[] = { argv[0] };
char** tmpargv = tmpargvv;
// glog“™‚Ì<E2809A>‰Šú‰»
caffe::GlobalInit(&tmpargc, &tmpargv);
});
boost::filesystem::path mode_dir_path(model_dir);
if (!mode_dir_path.is_absolute()) // model_dirが相対パスなら絶対パスに直す
{
@ -1287,6 +1327,8 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
void Waifu2x::destroy()
{
CudaDeviceSet devset(process, gpu_no);
net_noise.reset();
net_scale.reset();

4
common/waifu2x.h
View File

@ -117,6 +117,8 @@ private:
boost::optional<int> output_quality;
int output_depth;
int gpu_no;
private:
static eWaifu2xError LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file);
static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::vector<char> &img_data);
@ -162,7 +164,7 @@ public:
const boost::optional<double> scale_ratio, const boost::optional<int> scale_width, const boost::optional<int> scale_height,
const boost::filesystem::path &model_dir, const std::string &process,
const boost::optional<int> output_quality = boost::optional<int>(), const int output_depth = 8, const bool use_tta = false,
const int crop_size = 128, const int batch_size = 1);
const int crop_size = 128, const int batch_size = 1, const int gpu_no = 0);
void destroy();

57
waifu2x-caffe-gui/MainDialog.cpp
View File

@ -36,9 +36,21 @@ const TCHAR * const MultiFileStr = TEXT("(Multi File)");
const UINT_PTR nIDEventTimeLeft = 1000;
// http://stackoverflow.com/questions/10167382/boostfilesystem-get-relative-path
boost::filesystem::path relativePath(const boost::filesystem::path &path, const boost::filesystem::path &relative_to)
namespace
{
template<typename T>
static tstring to_tstring(T val)
{
#ifdef UNICODE
return std::to_wstring(val);
#else
return std::to_string(val);
#endif
}
// http://stackoverflow.com/questions/10167382/boostfilesystem-get-relative-path
boost::filesystem::path relativePath(const boost::filesystem::path &path, const boost::filesystem::path &relative_to)
{
// create absolute paths
boost::filesystem::path p = boost::filesystem::absolute(path);
boost::filesystem::path r = boost::filesystem::absolute(relative_to);
@ -78,10 +90,10 @@ boost::filesystem::path relativePath(const boost::filesystem::path &path, const
}
return result;
}
}
std::vector<int> CommonDivisorList(const int N)
{
std::vector<int> CommonDivisorList(const int N)
{
std::vector<int> list;
const int sq = sqrt(N);
@ -98,6 +110,7 @@ std::vector<int> CommonDivisorList(const int N)
std::sort(list.begin(), list.end());
return list;
}
}
@ -675,7 +688,7 @@ void DialogEvent::ProcessWaifu2x()
}
Waifu2x w;
ret = w.init(__argc, __argv, mode, noise_level, ScaleRatio, ScaleWidth, ScaleHeight, model_dir, process, output_quality, output_depth, use_tta, crop_size, batch_size);
ret = w.init(__argc, __argv, mode, noise_level, ScaleRatio, ScaleWidth, ScaleHeight, model_dir, process, output_quality, output_depth, use_tta, crop_size, batch_size, gpu_no);
if (ret != Waifu2x::eWaifu2xError_OK)
SendMessage(dh, WM_ON_WAIFU2X_ERROR, (WPARAM)&ret, 0);
else
@ -924,6 +937,8 @@ void DialogEvent::SaveIni(const bool isSyncMember)
WritePrivateProfileString(TEXT("Setting"), TEXT("LastUseTTA"), to_tstring(use_tta ? 1 : 0).c_str(), getTString(SettingFilePath).c_str());
WritePrivateProfileString(TEXT("Setting"), TEXT("LastDeviceNo"), to_tstring(gpu_no).c_str(), getTString(SettingFilePath).c_str());
if (output_quality)
WritePrivateProfileString(TEXT("Setting"), TEXT("LastOutputQuality"), boost::lexical_cast<tstring>(*output_quality).c_str(), getTString(SettingFilePath).c_str());
else
@ -1130,7 +1145,7 @@ UINT_PTR DialogEvent::OFNHookProcOut(HWND hdlg, UINT uiMsg, WPARAM wParam, LPARA
DialogEvent::DialogEvent() : dh(nullptr), mode("noise_scale"), noise_level(1), scale_ratio(2.0), scale_width(0), scale_height(0), model_dir(TEXT("models/anime_style_art_rgb")),
process("gpu"), outputExt(TEXT(".png")), inputFileExt(TEXT("png:jpg:jpeg:tif:tiff:bmp:tga")),
use_tta(false), output_depth(8), crop_size(128), batch_size(1), isLastError(false), scaleType(eScaleTypeEnd),
use_tta(false), output_depth(8), crop_size(128), batch_size(1), gpu_no(0), isLastError(false), scaleType(eScaleTypeEnd),
TimeLeftThread(-1), TimeLeftGetTimeThread(0), isCommandLineStart(false), tAutoMode(TEXT("none")),
isArgStartAuto(true), isArgStartSuccessFinish(true), isOutputNoOverwrite(false), isNotSaveParam(false)
{}
@ -1650,6 +1665,8 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
use_tta = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastUseTTA"), 0, getTString(SettingFilePath).c_str()) != 0;
gpu_no = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastDeviceNo"), 0, getTString(SettingFilePath).c_str());
output_quality.reset();
const int num = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastOutputQuality"), -100, getTString(SettingFilePath).c_str());
if (num != -100)
@ -1920,6 +1937,10 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
L"input batch size", false,
1, L"int", cmd);
TCLAP::ValueArg<int> cmdGPUNoFile(L"", L"gpu",
L"gpu device no", false,
0, L"int", cmd);
std::vector<int> cmdBoolConstraintV;
cmdBoolConstraintV.push_back(0);
cmdBoolConstraintV.push_back(1);
@ -2134,6 +2155,13 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
isSetParam = true;
}
if (cmdGPUNoFile.isSet())
{
gpu_no = cmdGPUNoFile.getValue();
isSetParam = true;
}
if (cmdTTA.isSet())
{
SendMessage(GetDlgItem(dh, IDC_CHECK_TTA), BM_SETCHECK, cmdTTA.getValue() ? BST_CHECKED : BST_UNCHECKED, 0);
@ -2844,6 +2872,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
bool isOutputNoOverwrite;
tstring tInputDirFix;
tstring tOutputDirFix;
int gpu_no;
private:
void AppSettingDialogEvent::SetWindowTextLang()
@ -2872,6 +2901,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
SET_WINDOW_TEXT(IDC_CHECK_OUTPUT_NO_OVERWIRITE);
SET_WINDOW_TEXT(IDC_STATIC_USE_GPU_NO);
#undef SET_WINDOW_TEXT
}
@ -2926,6 +2957,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
SetWindowText(GetDlgItem(hWnd, IDC_EDIT_INPUT_DIR_FIX), tInputDirFix.c_str());
SetWindowText(GetDlgItem(hWnd, IDC_EDIT_OUTPUT_DIR_FIX), tOutputDirFix.c_str());
SetWindowText(GetDlgItem(hWnd, IDC_EDIT_USE_GPU_NO), to_tstring(gpu_no).c_str());
}
void Close(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
@ -2935,6 +2968,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
void sync()
{
TCHAR *ptr = nullptr;
if (SendMessage(GetDlgItem(dh, IDC_RADIO_AUTO_START_ONE), BM_GETCHECK, 0, 0))
tAutoMode = TEXT("one");
else if (SendMessage(GetDlgItem(dh, IDC_RADIO_AUTO_START_MULTI), BM_GETCHECK, 0, 0))
@ -2960,6 +2995,12 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
GetWindowText(GetDlgItem(dh, IDC_EDIT_OUTPUT_DIR_FIX), buf, _countof(buf));
buf[_countof(buf) - 1] = TEXT('\0');
tOutputDirFix = buf;
GetWindowText(GetDlgItem(dh, IDC_EDIT_USE_GPU_NO), buf, _countof(buf));
buf[_countof(buf) - 1] = TEXT('\0');
const auto devno = _tcstol(buf, &ptr, 10);
if ((!ptr || *ptr == TEXT('\0')) && devno >= 0)
gpu_no = devno;
}
void AppSettingDialogEvent::OK(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
@ -2983,6 +3024,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
cAppSettingDialogEvent.isOutputNoOverwrite = isOutputNoOverwrite;
cAppSettingDialogEvent.tInputDirFix = tInputDirFix;
cAppSettingDialogEvent.tOutputDirFix = tOutputDirFix;
cAppSettingDialogEvent.gpu_no = gpu_no;
CDialog cDialog;
@ -3003,6 +3045,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
isOutputNoOverwrite = cAppSettingDialogEvent.isOutputNoOverwrite;
tInputDirFix = cAppSettingDialogEvent.tInputDirFix;
tOutputDirFix = cAppSettingDialogEvent.tOutputDirFix;
gpu_no = cAppSettingDialogEvent.gpu_no;
if (tOutputDirFix.length() > 0 && boost::filesystem::exists(tOutputDirFix))
{

12
waifu2x-caffe-gui/MainDialog.h
View File

@ -51,6 +51,8 @@ private:
int crop_size;
int batch_size;
int gpu_no;
std::vector<tstring> extList;
std::thread processThread;
@ -107,16 +109,6 @@ private:
bool isNotSaveParam;
private:
template<typename T>
static tstring to_tstring(T val)
{
#ifdef UNICODE
return std::to_wstring(val);
#else
return std::to_string(val);
#endif
}
tstring AddName() const;
bool SyncMember(const bool NotSyncCropSize, const bool silent = false);

BIN
waifu2x-caffe-gui/Resource.rc
View File

Binary file not shown.

BIN
waifu2x-caffe-gui/resource.h
View File

Binary file not shown.

6
waifu2x-caffe/Source.cpp
View File

@ -135,6 +135,10 @@ int main(int argc, char** argv)
"input batch size", false,
1, "int", cmd);
TCLAP::ValueArg<int> cmdGPUNoFile("", "gpu",
"gpu device no", false,
0, "int", cmd);
std::vector<int> cmdTTAConstraintV;
cmdTTAConstraintV.push_back(0);
cmdTTAConstraintV.push_back(1);
@ -319,7 +323,7 @@ int main(int argc, char** argv)
Waifu2x w;
ret = w.init(argc, argv, cmdMode.getValue(), cmdNRLevel.getValue(), ScaleRatio, ScaleWidth, ScaleHeight, cmdModelPath.getValue(), cmdProcess.getValue(),
cmdOutputQuality.getValue() == -1 ? boost::optional<int>() : cmdOutputQuality.getValue(), cmdOutputDepth.getValue(), use_tta, cmdCropSizeFile.getValue(),
cmdBatchSizeFile.getValue());
cmdBatchSizeFile.getValue(), cmdGPUNoFile.getValue());
switch (ret)
{
case Waifu2x::eWaifu2xError_InvalidParameter: