benchmark.lua 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672
  1. require 'pl'
  2. local __FILE__ = (function() return string.gsub(debug.getinfo(2, 'S').source, "^@", "") end)()
  3. package.path = path.join(path.dirname(__FILE__), "..", "lib", "?.lua;") .. package.path
  4. require 'xlua'
  5. require 'w2nn'
  6. local iproc = require 'iproc'
  7. local reconstruct = require 'reconstruct'
  8. local image_loader = require 'image_loader'
  9. local gm = require 'graphicsmagick'
  10. local cjson = require 'cjson'
  11. local cmd = torch.CmdLine()
  12. cmd:text()
  13. cmd:text("waifu2x-benchmark")
  14. cmd:text("Options:")
  15. cmd:option("-dir", "./data/test", 'test image directory')
  16. cmd:option("-file", "", 'test image file list')
  17. cmd:option("-model1_dir", "./models/anime_style_art_rgb", 'model1 directory')
  18. cmd:option("-model2_dir", "", 'model2 directory (optional)')
  19. cmd:option("-method", "scale", '(scale|noise|noise_scale|user|diff)')
  20. cmd:option("-filter", "Catrom", "downscaling filter (Box|Lanczos|Catrom(Bicubic))")
  21. cmd:option("-resize_blur", 1.0, 'blur parameter for resize')
  22. cmd:option("-color", "y", '(rgb|y|r|g|b)')
  23. cmd:option("-noise_level", 1, 'model noise level')
  24. cmd:option("-jpeg_quality", 75, 'jpeg quality')
  25. cmd:option("-jpeg_times", 1, 'jpeg compression times')
  26. cmd:option("-jpeg_quality_down", 5, 'value of jpeg quality to decrease each times')
  27. cmd:option("-range_bug", 0, 'Reproducing the dynamic range bug that is caused by MATLAB\'s rgb2ycbcr(1|0)')
  28. cmd:option("-save_image", 0, 'save converted images')
  29. cmd:option("-save_baseline_image", 0, 'save baseline images')
  30. cmd:option("-output_dir", "./", 'output directroy')
  31. cmd:option("-show_progress", 1, 'show progressbar')
  32. cmd:option("-baseline_filter", "Catrom", 'baseline interpolation (Box|Lanczos|Catrom(Bicubic))')
  33. cmd:option("-save_info", 0, 'save score and parameters to benchmark.txt')
  34. cmd:option("-save_all", 0, 'group -save_info, -save_image and -save_baseline_image option')
  35. cmd:option("-thread", -1, 'number of CPU threads')
  36. cmd:option("-tta", 0, 'use tta')
  37. cmd:option("-tta_level", 8, 'tta level')
  38. cmd:option("-crop_size", 128, 'patch size per process')
  39. cmd:option("-batch_size", 1, 'batch_size')
  40. cmd:option("-force_cudnn", 0, 'use cuDNN backend')
  41. cmd:option("-yuv420", 0, 'use yuv420 jpeg')
  42. cmd:option("-name", "", 'model name for user method')
  43. cmd:option("-x_dir", "", 'input image for user method')
  44. cmd:option("-y_dir", "", 'groundtruth image for user method. filename must be the same as x_dir')
  45. cmd:option("-x_file", "", 'input image for user method')
  46. cmd:option("-y_file", "", 'groundtruth image for user method. filename must be the same as x_file')
  47. cmd:option("-border", 0, 'border px that will removed')
  48. local function to_bool(settings, name)
  49. if settings[name] == 1 then
  50. settings[name] = true
  51. else
  52. settings[name] = false
  53. end
  54. end
  55. local opt = cmd:parse(arg)
  56. torch.setdefaulttensortype('torch.FloatTensor')
  57. if cudnn then
  58. cudnn.fastest = true
  59. cudnn.benchmark = true
  60. end
  61. to_bool(opt, "force_cudnn")
  62. to_bool(opt, "yuv420")
  63. to_bool(opt, "save_all")
  64. to_bool(opt, "tta")
  65. if opt.save_all then
  66. opt.save_image = true
  67. opt.save_info = true
  68. opt.save_baseline_image = true
  69. else
  70. to_bool(opt, "save_image")
  71. to_bool(opt, "save_info")
  72. to_bool(opt, "save_baseline_image")
  73. end
  74. to_bool(opt, "show_progress")
  75. if opt.thread > 0 then
  76. torch.setnumthreads(tonumber(opt.thread))
  77. end
  78. if opt.output_dir:len() > 0 then
  79. dir.makepath(opt.output_dir)
  80. end
  81. local function rgb2y_matlab(x)
  82. local y = torch.Tensor(1, x:size(2), x:size(3)):zero()
  83. x = iproc.byte2float(x)
  84. y:add(x[1] * 65.481)
  85. y:add(x[2] * 128.553)
  86. y:add(x[3] * 24.966)
  87. y:add(16.0)
  88. return y:byte():float()
  89. end
  90. local function RGBMSE(x1, x2)
  91. x1 = iproc.float2byte(x1):float()
  92. x2 = iproc.float2byte(x2):float()
  93. return (x1 - x2):pow(2):mean()
  94. end
  95. local function CHMSE(x1, x2, ch)
  96. x1 = iproc.float2byte(x1):float()
  97. x2 = iproc.float2byte(x2):float()
  98. return (x1[ch] - x2[ch]):pow(2):mean()
  99. end
  100. local function YMSE(x1, x2)
  101. if opt.range_bug == 1 then
  102. local x1_2 = rgb2y_matlab(x1)
  103. local x2_2 = rgb2y_matlab(x2)
  104. return (x1_2 - x2_2):pow(2):mean()
  105. else
  106. local x1_2 = image.rgb2y(x1):mul(255.0)
  107. local x2_2 = image.rgb2y(x2):mul(255.0)
  108. return (x1_2 - x2_2):pow(2):mean()
  109. end
  110. end
  111. local function MSE(x1, x2, color)
  112. if color == "y" then
  113. return YMSE(x1, x2)
  114. elseif color == "r" then
  115. return CHMSE(x1, x2, 1)
  116. elseif color == "g" then
  117. return CHMSE(x1, x2, 2)
  118. elseif color == "b" then
  119. return CHMSE(x1, x2, 3)
  120. else
  121. return RGBMSE(x1, x2)
  122. end
  123. end
  124. local function PSNR(x1, x2, color)
  125. local mse = math.max(MSE(x1, x2, color), 1)
  126. return 10 * math.log10((255.0 * 255.0) / mse)
  127. end
  128. local function MSE2PSNR(mse)
  129. return 10 * math.log10((255.0 * 255.0) / math.max(mse, 1))
  130. end
  131. local function transform_jpeg(x, opt)
  132. for i = 1, opt.jpeg_times do
  133. jpeg = gm.Image(x, "RGB", "DHW")
  134. jpeg:format("jpeg")
  135. if opt.yuv420 then
  136. jpeg:samplingFactors({2.0, 1.0, 1.0})
  137. else
  138. jpeg:samplingFactors({1.0, 1.0, 1.0})
  139. end
  140. blob, len = jpeg:toBlob(opt.jpeg_quality - (i - 1) * opt.jpeg_quality_down)
  141. jpeg:fromBlob(blob, len)
  142. x = jpeg:toTensor("byte", "RGB", "DHW")
  143. end
  144. return iproc.byte2float(x)
  145. end
  146. local function baseline_scale(x, filter)
  147. return iproc.scale(x,
  148. x:size(3) * 2.0,
  149. x:size(2) * 2.0,
  150. filter)
  151. end
  152. local function transform_scale(x, opt)
  153. return iproc.scale(x,
  154. x:size(3) * 0.5,
  155. x:size(2) * 0.5,
  156. opt.filter, opt.resize_blur)
  157. end
  158. local function transform_scale_jpeg(x, opt)
  159. x = iproc.scale(x,
  160. x:size(3) * 0.5,
  161. x:size(2) * 0.5,
  162. opt.filter, opt.resize_blur)
  163. for i = 1, opt.jpeg_times do
  164. jpeg = gm.Image(x, "RGB", "DHW")
  165. jpeg:format("jpeg")
  166. if opt.yuv420 then
  167. jpeg:samplingFactors({2.0, 1.0, 1.0})
  168. else
  169. jpeg:samplingFactors({1.0, 1.0, 1.0})
  170. end
  171. blob, len = jpeg:toBlob(opt.jpeg_quality - (i - 1) * opt.jpeg_quality_down)
  172. jpeg:fromBlob(blob, len)
  173. x = jpeg:toTensor("byte", "RGB", "DHW")
  174. end
  175. return iproc.byte2float(x)
  176. end
  177. local function remove_border(x, border)
  178. return iproc.crop(x,
  179. border, border,
  180. x:size(3) - border,
  181. x:size(2) - border)
  182. end
  183. local function benchmark(opt, x, model1, model2)
  184. local mse1, mse2
  185. local won = {0, 0}
  186. local model1_mse = 0
  187. local model2_mse = 0
  188. local baseline_mse = 0
  189. local model1_psnr = 0
  190. local model2_psnr = 0
  191. local baseline_psnr = 0
  192. local model1_time = 0
  193. local model2_time = 0
  194. local scale_f = reconstruct.scale
  195. local image_f = reconstruct.image
  196. local detail_fp = nil
  197. if opt.save_info then
  198. detail_fp = io.open(path.join(opt.output_dir, "benchmark_details.txt"), "w")
  199. end
  200. if opt.tta then
  201. scale_f = function(model, scale, x, block_size, batch_size)
  202. return reconstruct.scale_tta(model, opt.tta_level,
  203. scale, x, block_size, batch_size)
  204. end
  205. image_f = function(model, x, block_size, batch_size)
  206. return reconstruct.image_tta(model, opt.tta_level,
  207. x, block_size, batch_size)
  208. end
  209. end
  210. for i = 1, #x do
  211. local basename = x[i].basename
  212. local input, model1_output, model2_output, baseline_output, ground_truth
  213. if opt.method == "scale" then
  214. input = transform_scale(x[i].y, opt)
  215. ground_truth = x[i].y
  216. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  217. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  218. if model2 then
  219. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  220. end
  221. end
  222. t = sys.clock()
  223. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  224. model1_time = model1_time + (sys.clock() - t)
  225. if model2 then
  226. t = sys.clock()
  227. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  228. model2_time = model2_time + (sys.clock() - t)
  229. end
  230. baseline_output = baseline_scale(input, opt.baseline_filter)
  231. elseif opt.method == "noise" then
  232. input = transform_jpeg(x[i].y, opt)
  233. ground_truth = x[i].y
  234. if opt.force_cudnn and i == 1 then
  235. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  236. if model2 then
  237. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  238. end
  239. end
  240. t = sys.clock()
  241. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  242. model1_time = model1_time + (sys.clock() - t)
  243. if model2 then
  244. t = sys.clock()
  245. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  246. model2_time = model2_time + (sys.clock() - t)
  247. end
  248. baseline_output = input
  249. elseif opt.method == "noise_scale" then
  250. input = transform_scale_jpeg(x[i].y, opt)
  251. ground_truth = x[i].y
  252. if opt.force_cudnn and i == 1 then
  253. if model1.noise_scale_model then
  254. model1_output = scale_f(model1.noise_scale_model, 2.0,
  255. input, opt.crop_size, opt.batch_size)
  256. else
  257. if model1.noise_model then
  258. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  259. else
  260. model1_output = input
  261. end
  262. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  263. opt.crop_size, opt.batch_size)
  264. end
  265. if model2 then
  266. if model2.noise_scale_model then
  267. model2_output = scale_f(model2.noise_scale_model, 2.0,
  268. input, opt.crop_size, opt.batch_size)
  269. else
  270. if model2.noise_model then
  271. model2_output = image_f(model2.noise_model, input,
  272. opt.crop_size, opt.batch_size)
  273. else
  274. model2_output = input
  275. end
  276. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  277. opt.crop_size, opt.batch_size)
  278. end
  279. end
  280. end
  281. t = sys.clock()
  282. if model1.noise_scale_model then
  283. model1_output = scale_f(model1.noise_scale_model, 2.0,
  284. input, opt.crop_size, opt.batch_size)
  285. else
  286. if model1.noise_model then
  287. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  288. else
  289. model1_output = input
  290. end
  291. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  292. opt.crop_size, opt.batch_size)
  293. end
  294. model1_time = model1_time + (sys.clock() - t)
  295. if model2 then
  296. t = sys.clock()
  297. if model2.noise_scale_model then
  298. model2_output = scale_f(model2.noise_scale_model, 2.0,
  299. input, opt.crop_size, opt.batch_size)
  300. else
  301. if model2.noise_model then
  302. model2_output = image_f(model2.noise_model, input,
  303. opt.crop_size, opt.batch_size)
  304. else
  305. model2_output = input
  306. end
  307. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  308. opt.crop_size, opt.batch_size)
  309. end
  310. model2_time = model2_time + (sys.clock() - t)
  311. end
  312. baseline_output = baseline_scale(input, opt.baseline_filter)
  313. elseif opt.method == "user" then
  314. input = x[i].x
  315. ground_truth = x[i].y
  316. local y_scale = ground_truth:size(2) / input:size(2)
  317. if y_scale > 1 then
  318. if opt.force_cudnn and i == 1 then
  319. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  320. if model2 then
  321. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  322. end
  323. end
  324. t = sys.clock()
  325. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  326. model1_time = model1_time + (sys.clock() - t)
  327. if model2 then
  328. t = sys.clock()
  329. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  330. model2_time = model2_time + (sys.clock() - t)
  331. end
  332. else
  333. if opt.force_cudnn and i == 1 then
  334. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  335. if model2 then
  336. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  337. end
  338. end
  339. t = sys.clock()
  340. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  341. model1_time = model1_time + (sys.clock() - t)
  342. if model2 then
  343. t = sys.clock()
  344. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  345. model2_time = model2_time + (sys.clock() - t)
  346. end
  347. end
  348. elseif opt.method == "diff" then
  349. input = x[i].x
  350. ground_truth = x[i].y
  351. model1_output = input
  352. end
  353. if opt.border > 0 then
  354. ground_truth = remove_border(ground_truth, opt.border)
  355. model1_output = remove_border(model1_output, opt.border)
  356. end
  357. mse1 = MSE(ground_truth, model1_output, opt.color)
  358. model1_mse = model1_mse + mse1
  359. model1_psnr = model1_psnr + MSE2PSNR(mse1)
  360. local won_model = 1
  361. if model2 then
  362. if opt.border > 0 then
  363. model2_output = remove_border(model2_output, opt.border)
  364. end
  365. mse2 = MSE(ground_truth, model2_output, opt.color)
  366. model2_mse = model2_mse + mse2
  367. model2_psnr = model2_psnr + MSE2PSNR(mse2)
  368. if mse1 < mse2 then
  369. won[1] = won[1] + 1
  370. elseif mse1 > mse2 then
  371. won[2] = won[2] + 1
  372. won_model = 2
  373. end
  374. if detail_fp then
  375. detail_fp:write(string.format("%s,%f,%f,%d\n", x[i].basename,
  376. MSE2PSNR(mse1), MSE2PSNR(mse2), won_model))
  377. end
  378. else
  379. if detail_fp then
  380. detail_fp:write(string.format("%s,%f\n", x[i].basename, MSE2PSNR(mse1)))
  381. end
  382. end
  383. if baseline_output then
  384. baseline_output = remove_border(baseline_output, opt.border)
  385. mse = MSE(ground_truth, baseline_output, opt.color)
  386. baseline_mse = baseline_mse + mse
  387. baseline_psnr = baseline_psnr + MSE2PSNR(mse)
  388. end
  389. if opt.save_image then
  390. if opt.save_baseline_image and baseline_output then
  391. image.save(path.join(opt.output_dir, string.format("%s_baseline.png", basename)),
  392. baseline_output)
  393. end
  394. if model1_output then
  395. image.save(path.join(opt.output_dir, string.format("%s_model1.png", basename)),
  396. model1_output)
  397. end
  398. if model2_output then
  399. image.save(path.join(opt.output_dir, string.format("%s_model2.png", basename)),
  400. model2_output)
  401. end
  402. end
  403. if opt.show_progress or i == #x then
  404. if model2 then
  405. if baseline_output then
  406. io.stdout:write(
  407. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, model2_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f, model2_psnr=%.3f, model1_won=%d, model2_won=%d \r",
  408. i, #x,
  409. model1_time,
  410. model2_time,
  411. math.sqrt(baseline_mse / i),
  412. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  413. baseline_psnr / i,
  414. model1_psnr / i, model2_psnr / i,
  415. won[1], won[2]
  416. ))
  417. else
  418. io.stdout:write(
  419. string.format("%d/%d; model1_time=%.2f, model2_time=%.2f, model1_rmse=%.3f, model2_rmse=%.3f, model1_psnr=%.3f, model2_psnr=%.3f, model1_own=%d, model2_won=%d \r",
  420. i, #x,
  421. model1_time,
  422. model2_time,
  423. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  424. model1_psnr / i, model2_psnr / i,
  425. won[1], won[2]
  426. ))
  427. end
  428. else
  429. if baseline_output then
  430. io.stdout:write(
  431. string.format("%d/%d; model1_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f \r",
  432. i, #x,
  433. model1_time,
  434. math.sqrt(baseline_mse / i), math.sqrt(model1_mse / i),
  435. baseline_psnr / i, model1_psnr / i
  436. ))
  437. else
  438. io.stdout:write(
  439. string.format("%d/%d; model1_time=%.2f, model1_rmse=%.3f, model1_psnr=%.3f \r",
  440. i, #x,
  441. model1_time,
  442. math.sqrt(model1_mse / i), model1_psnr / i
  443. ))
  444. end
  445. end
  446. io.stdout:flush()
  447. end
  448. end
  449. if opt.save_info then
  450. local fp = io.open(path.join(opt.output_dir, "benchmark.txt"), "w")
  451. fp:write("options : " .. cjson.encode(opt) .. "\n")
  452. if baseline_psnr > 0 then
  453. fp:write(string.format("baseline: RMSE = %.3f, PSNR = %.3f\n",
  454. math.sqrt(baseline_mse / #x), baseline_psnr / #x))
  455. end
  456. if model1_psnr > 0 then
  457. fp:write(string.format("model1 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  458. math.sqrt(model1_mse / #x), model1_psnr / #x, model1_time))
  459. end
  460. if model2_psnr > 0 then
  461. fp:write(string.format("model2 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  462. math.sqrt(model2_mse / #x), model2_psnr / #x, model2_time))
  463. end
  464. fp:close()
  465. if detail_fp then
  466. detail_fp:close()
  467. end
  468. end
  469. io.stdout:write("\n")
  470. end
  471. local function load_data_from_dir(test_dir)
  472. local test_x = {}
  473. local files = dir.getfiles(test_dir, "*.*")
  474. for i = 1, #files do
  475. local name = path.basename(files[i])
  476. local e = path.extension(name)
  477. local base = name:sub(0, name:len() - e:len())
  478. local img = image_loader.load_float(files[i])
  479. if img then
  480. table.insert(test_x, {y = iproc.crop_mod4(img),
  481. basename = base})
  482. end
  483. if opt.show_progress then
  484. xlua.progress(i, #files)
  485. end
  486. end
  487. return test_x
  488. end
  489. local function load_data_from_file(test_file)
  490. local test_x = {}
  491. local files = utils.split(file.read(test_file), "\n")
  492. for i = 1, #files do
  493. local name = path.basename(files[i])
  494. local e = path.extension(name)
  495. local base = name:sub(0, name:len() - e:len())
  496. local img = image_loader.load_float(files[i])
  497. if img then
  498. table.insert(test_x, {y = iproc.crop_mod4(img),
  499. basename = base})
  500. end
  501. if opt.show_progress then
  502. xlua.progress(i, #files)
  503. end
  504. end
  505. return test_x
  506. end
  507. local function get_basename(f)
  508. local name = path.basename(f)
  509. local e = path.extension(name)
  510. local base = name:sub(0, name:len() - e:len())
  511. return base
  512. end
  513. local function load_user_data(y_dir, y_file, x_dir, x_file)
  514. local test = {}
  515. local y_files
  516. local x_files
  517. if y_file:len() > 0 then
  518. y_files = utils.split(file.read(y_file), "\n")
  519. else
  520. y_files = dir.getfiles(y_dir, "*.*")
  521. end
  522. if x_file:len() > 0 then
  523. x_files = utils.split(file.read(x_file), "\n")
  524. else
  525. x_files = dir.getfiles(x_dir, "*.*")
  526. end
  527. local basename_db = {}
  528. for i = 1, #y_files do
  529. basename_db[get_basename(y_files[i])] = {y = y_files[i]}
  530. end
  531. for i = 1, #x_files do
  532. local key = get_basename(x_files[i])
  533. if basename_db[key] then
  534. basename_db[key].x = x_files[i]
  535. else
  536. error(string.format("%s is not found in %s", key, y_dir))
  537. end
  538. end
  539. for i = 1, #y_files do
  540. local key = get_basename(y_files[i])
  541. local d = basename_db[key]
  542. if not (d.x and d.y) then
  543. error(string.format("%s is not found in %s", key, x_dir))
  544. end
  545. end
  546. for i = 1, #y_files do
  547. local key = get_basename(y_files[i])
  548. local x = image_loader.load_float(basename_db[key].x)
  549. local y = image_loader.load_float(basename_db[key].y)
  550. if x and y then
  551. table.insert(test, {y = y,
  552. x = x,
  553. basename = key})
  554. end
  555. if opt.show_progress then
  556. xlua.progress(i, #y_files)
  557. end
  558. end
  559. return test
  560. end
  561. function load_noise_scale_model(model_dir, noise_level, force_cudnn)
  562. local f = path.join(model_dir, string.format("noise%d_scale2.0x_model.t7", opt.noise_level))
  563. local s1, noise_scale = pcall(w2nn.load_model, f, force_cudnn)
  564. local model = {}
  565. if not s1 then
  566. f = path.join(model_dir, string.format("noise%d_model.t7", opt.noise_level))
  567. local noise
  568. s1, noise = pcall(w2nn.load_model, f, force_cudnn)
  569. if not s1 then
  570. model.noise_model = nil
  571. print(model_dir .. "'s noise model is not found. benchmark will use only scale model.")
  572. else
  573. model.noise_model = noise
  574. end
  575. f = path.join(model_dir, "scale2.0x_model.t7")
  576. local scale
  577. s1, scale = pcall(w2nn.load_model, f, force_cudnn)
  578. if not s1 then
  579. error(model_dir .. ": load error")
  580. return nil
  581. end
  582. model.scale_model = scale
  583. else
  584. model.noise_scale_model = noise_scale
  585. end
  586. return model
  587. end
  588. if opt.show_progress then
  589. print(opt)
  590. end
  591. if opt.method == "scale" then
  592. local f1 = path.join(opt.model1_dir, "scale2.0x_model.t7")
  593. local f2 = path.join(opt.model2_dir, "scale2.0x_model.t7")
  594. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  595. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  596. if not s1 then
  597. error("Load error: " .. f1)
  598. end
  599. if not s2 then
  600. model2 = nil
  601. end
  602. local test_x
  603. if opt.file:len() > 0 then
  604. test_x = load_data_from_file(opt.file)
  605. else
  606. test_x = load_data_from_dir(opt.dir)
  607. end
  608. benchmark(opt, test_x, model1, model2)
  609. elseif opt.method == "noise" then
  610. local f1 = path.join(opt.model1_dir, string.format("noise%d_model.t7", opt.noise_level))
  611. local f2 = path.join(opt.model2_dir, string.format("noise%d_model.t7", opt.noise_level))
  612. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  613. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  614. if not s1 then
  615. error("Load error: " .. f1)
  616. end
  617. if not s2 then
  618. model2 = nil
  619. end
  620. local test_x
  621. if opt.file:len() > 0 then
  622. test_x = load_data_from_file(opt.file)
  623. else
  624. test_x = load_data_from_dir(opt.dir)
  625. end
  626. benchmark(opt, test_x, model1, model2)
  627. elseif opt.method == "noise_scale" then
  628. local model2 = nil
  629. local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level, opt.force_cudnn)
  630. if opt.model2_dir:len() > 0 then
  631. model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level, opt.force_cudnn)
  632. end
  633. local test_x
  634. if opt.file:len() > 0 then
  635. test_x = load_data_from_file(opt.file)
  636. else
  637. test_x = load_data_from_dir(opt.dir)
  638. end
  639. benchmark(opt, test_x, model1, model2)
  640. elseif opt.method == "user" then
  641. local f1 = path.join(opt.model1_dir, string.format("%s_model.t7", opt.name))
  642. local f2 = path.join(opt.model2_dir, string.format("%s_model.t7", opt.name))
  643. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  644. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  645. if not s1 then
  646. error("Load error: " .. f1)
  647. end
  648. if not s2 then
  649. model2 = nil
  650. end
  651. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  652. benchmark(opt, test, model1, model2)
  653. elseif opt.method == "diff" then
  654. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  655. benchmark(opt, test, nil, nil)
  656. end