benchmark.lua 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704
  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|scale4)')
  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 baseline_scale4(x, filter)
  153. return iproc.scale(x,
  154. x:size(3) * 4.0,
  155. x:size(2) * 4.0,
  156. filter)
  157. end
  158. local function transform_scale(x, opt)
  159. return iproc.scale(x,
  160. x:size(3) * 0.5,
  161. x:size(2) * 0.5,
  162. opt.filter, opt.resize_blur)
  163. end
  164. local function transform_scale4(x, opt)
  165. return iproc.scale(x,
  166. x:size(3) * 0.25,
  167. x:size(2) * 0.25,
  168. opt.filter, opt.resize_blur)
  169. end
  170. local function transform_scale_jpeg(x, opt)
  171. x = iproc.scale(x,
  172. x:size(3) * 0.5,
  173. x:size(2) * 0.5,
  174. opt.filter, opt.resize_blur)
  175. for i = 1, opt.jpeg_times do
  176. jpeg = gm.Image(x, "RGB", "DHW")
  177. jpeg:format("jpeg")
  178. if opt.yuv420 then
  179. jpeg:samplingFactors({2.0, 1.0, 1.0})
  180. else
  181. jpeg:samplingFactors({1.0, 1.0, 1.0})
  182. end
  183. blob, len = jpeg:toBlob(opt.jpeg_quality - (i - 1) * opt.jpeg_quality_down)
  184. jpeg:fromBlob(blob, len)
  185. x = jpeg:toTensor("byte", "RGB", "DHW")
  186. end
  187. return iproc.byte2float(x)
  188. end
  189. local function remove_border(x, border)
  190. return iproc.crop(x,
  191. border, border,
  192. x:size(3) - border,
  193. x:size(2) - border)
  194. end
  195. local function benchmark(opt, x, model1, model2)
  196. local mse1, mse2
  197. local won = {0, 0}
  198. local model1_mse = 0
  199. local model2_mse = 0
  200. local baseline_mse = 0
  201. local model1_psnr = 0
  202. local model2_psnr = 0
  203. local baseline_psnr = 0
  204. local model1_time = 0
  205. local model2_time = 0
  206. local scale_f = reconstruct.scale
  207. local image_f = reconstruct.image
  208. local detail_fp = nil
  209. if opt.save_info then
  210. detail_fp = io.open(path.join(opt.output_dir, "benchmark_details.txt"), "w")
  211. end
  212. if opt.tta then
  213. scale_f = function(model, scale, x, block_size, batch_size)
  214. return reconstruct.scale_tta(model, opt.tta_level,
  215. scale, x, block_size, batch_size)
  216. end
  217. image_f = function(model, x, block_size, batch_size)
  218. return reconstruct.image_tta(model, opt.tta_level,
  219. x, block_size, batch_size)
  220. end
  221. end
  222. for i = 1, #x do
  223. local basename = x[i].basename
  224. local input, model1_output, model2_output, baseline_output, ground_truth
  225. if opt.method == "scale" then
  226. input = transform_scale(x[i].y, opt)
  227. ground_truth = x[i].y
  228. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  229. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  230. if model2 then
  231. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  232. end
  233. end
  234. t = sys.clock()
  235. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  236. model1_time = model1_time + (sys.clock() - t)
  237. if model2 then
  238. t = sys.clock()
  239. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  240. model2_time = model2_time + (sys.clock() - t)
  241. end
  242. baseline_output = baseline_scale(input, opt.baseline_filter)
  243. elseif opt.method == "scale4" then
  244. input = transform_scale4(x[i].y, opt)
  245. ground_truth = x[i].y
  246. if opt.force_cudnn and i == 1 then -- run cuDNN benchmark first
  247. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  248. if model2 then
  249. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  250. end
  251. end
  252. t = sys.clock()
  253. model1_output = scale_f(model1, 2.0, input, opt.crop_size, opt.batch_size)
  254. model1_output = scale_f(model1, 2.0, model1_output, opt.crop_size, opt.batch_size)
  255. model1_time = model1_time + (sys.clock() - t)
  256. if model2 then
  257. t = sys.clock()
  258. model2_output = scale_f(model2, 2.0, input, opt.crop_size, opt.batch_size)
  259. model2_output = scale_f(model2, 2.0, model2_output, opt.crop_size, opt.batch_size)
  260. model2_time = model2_time + (sys.clock() - t)
  261. end
  262. baseline_output = baseline_scale4(input, opt.baseline_filter)
  263. elseif opt.method == "noise" then
  264. input = transform_jpeg(x[i].y, opt)
  265. ground_truth = x[i].y
  266. if opt.force_cudnn and i == 1 then
  267. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  268. if model2 then
  269. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  270. end
  271. end
  272. t = sys.clock()
  273. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  274. model1_time = model1_time + (sys.clock() - t)
  275. if model2 then
  276. t = sys.clock()
  277. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  278. model2_time = model2_time + (sys.clock() - t)
  279. end
  280. baseline_output = input
  281. elseif opt.method == "noise_scale" then
  282. input = transform_scale_jpeg(x[i].y, opt)
  283. ground_truth = x[i].y
  284. if opt.force_cudnn and i == 1 then
  285. if model1.noise_scale_model then
  286. model1_output = scale_f(model1.noise_scale_model, 2.0,
  287. input, opt.crop_size, opt.batch_size)
  288. else
  289. if model1.noise_model then
  290. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  291. else
  292. model1_output = input
  293. end
  294. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  295. opt.crop_size, opt.batch_size)
  296. end
  297. if model2 then
  298. if model2.noise_scale_model then
  299. model2_output = scale_f(model2.noise_scale_model, 2.0,
  300. input, opt.crop_size, opt.batch_size)
  301. else
  302. if model2.noise_model then
  303. model2_output = image_f(model2.noise_model, input,
  304. opt.crop_size, opt.batch_size)
  305. else
  306. model2_output = input
  307. end
  308. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  309. opt.crop_size, opt.batch_size)
  310. end
  311. end
  312. end
  313. t = sys.clock()
  314. if model1.noise_scale_model then
  315. model1_output = scale_f(model1.noise_scale_model, 2.0,
  316. input, opt.crop_size, opt.batch_size)
  317. else
  318. if model1.noise_model then
  319. model1_output = image_f(model1.noise_model, input, opt.crop_size, opt.batch_size)
  320. else
  321. model1_output = input
  322. end
  323. model1_output = scale_f(model1.scale_model, 2.0, model1_output,
  324. opt.crop_size, opt.batch_size)
  325. end
  326. model1_time = model1_time + (sys.clock() - t)
  327. if model2 then
  328. t = sys.clock()
  329. if model2.noise_scale_model then
  330. model2_output = scale_f(model2.noise_scale_model, 2.0,
  331. input, opt.crop_size, opt.batch_size)
  332. else
  333. if model2.noise_model then
  334. model2_output = image_f(model2.noise_model, input,
  335. opt.crop_size, opt.batch_size)
  336. else
  337. model2_output = input
  338. end
  339. model2_output = scale_f(model2.scale_model, 2.0, model2_output,
  340. opt.crop_size, opt.batch_size)
  341. end
  342. model2_time = model2_time + (sys.clock() - t)
  343. end
  344. baseline_output = baseline_scale(input, opt.baseline_filter)
  345. elseif opt.method == "user" then
  346. input = x[i].x
  347. ground_truth = x[i].y
  348. local y_scale = ground_truth:size(2) / input:size(2)
  349. if y_scale > 1 then
  350. if opt.force_cudnn and i == 1 then
  351. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  352. if model2 then
  353. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  354. end
  355. end
  356. t = sys.clock()
  357. model1_output = scale_f(model1, y_scale, input, opt.crop_size, opt.batch_size)
  358. model1_time = model1_time + (sys.clock() - t)
  359. if model2 then
  360. t = sys.clock()
  361. model2_output = scale_f(model2, y_scale, input, opt.crop_size, opt.batch_size)
  362. model2_time = model2_time + (sys.clock() - t)
  363. end
  364. else
  365. if opt.force_cudnn and i == 1 then
  366. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  367. if model2 then
  368. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  369. end
  370. end
  371. t = sys.clock()
  372. model1_output = image_f(model1, input, opt.crop_size, opt.batch_size)
  373. model1_time = model1_time + (sys.clock() - t)
  374. if model2 then
  375. t = sys.clock()
  376. model2_output = image_f(model2, input, opt.crop_size, opt.batch_size)
  377. model2_time = model2_time + (sys.clock() - t)
  378. end
  379. end
  380. elseif opt.method == "diff" then
  381. input = x[i].x
  382. ground_truth = x[i].y
  383. model1_output = input
  384. end
  385. if opt.border > 0 then
  386. ground_truth = remove_border(ground_truth, opt.border)
  387. model1_output = remove_border(model1_output, opt.border)
  388. end
  389. mse1 = MSE(ground_truth, model1_output, opt.color)
  390. model1_mse = model1_mse + mse1
  391. model1_psnr = model1_psnr + MSE2PSNR(mse1)
  392. local won_model = 1
  393. if model2 then
  394. if opt.border > 0 then
  395. model2_output = remove_border(model2_output, opt.border)
  396. end
  397. mse2 = MSE(ground_truth, model2_output, opt.color)
  398. model2_mse = model2_mse + mse2
  399. model2_psnr = model2_psnr + MSE2PSNR(mse2)
  400. if mse1 < mse2 then
  401. won[1] = won[1] + 1
  402. elseif mse1 > mse2 then
  403. won[2] = won[2] + 1
  404. won_model = 2
  405. end
  406. if detail_fp then
  407. detail_fp:write(string.format("%s,%f,%f,%d\n", x[i].basename,
  408. MSE2PSNR(mse1), MSE2PSNR(mse2), won_model))
  409. end
  410. else
  411. if detail_fp then
  412. detail_fp:write(string.format("%s,%f\n", x[i].basename, MSE2PSNR(mse1)))
  413. end
  414. end
  415. if baseline_output then
  416. baseline_output = remove_border(baseline_output, opt.border)
  417. mse = MSE(ground_truth, baseline_output, opt.color)
  418. baseline_mse = baseline_mse + mse
  419. baseline_psnr = baseline_psnr + MSE2PSNR(mse)
  420. end
  421. if opt.save_image then
  422. if opt.save_baseline_image and baseline_output then
  423. image.save(path.join(opt.output_dir, string.format("%s_baseline.png", basename)),
  424. baseline_output)
  425. end
  426. if model1_output then
  427. image.save(path.join(opt.output_dir, string.format("%s_model1.png", basename)),
  428. model1_output)
  429. end
  430. if model2_output then
  431. image.save(path.join(opt.output_dir, string.format("%s_model2.png", basename)),
  432. model2_output)
  433. end
  434. end
  435. if opt.show_progress or i == #x then
  436. if model2 then
  437. if baseline_output then
  438. io.stdout:write(
  439. 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",
  440. i, #x,
  441. model1_time,
  442. model2_time,
  443. math.sqrt(baseline_mse / i),
  444. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  445. baseline_psnr / i,
  446. model1_psnr / i, model2_psnr / i,
  447. won[1], won[2]
  448. ))
  449. else
  450. io.stdout:write(
  451. 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",
  452. i, #x,
  453. model1_time,
  454. model2_time,
  455. math.sqrt(model1_mse / i), math.sqrt(model2_mse / i),
  456. model1_psnr / i, model2_psnr / i,
  457. won[1], won[2]
  458. ))
  459. end
  460. else
  461. if baseline_output then
  462. io.stdout:write(
  463. string.format("%d/%d; model1_time=%.2f, baseline_rmse=%.3f, model1_rmse=%.3f, baseline_psnr=%.3f, model1_psnr=%.3f \r",
  464. i, #x,
  465. model1_time,
  466. math.sqrt(baseline_mse / i), math.sqrt(model1_mse / i),
  467. baseline_psnr / i, model1_psnr / i
  468. ))
  469. else
  470. io.stdout:write(
  471. string.format("%d/%d; model1_time=%.2f, model1_rmse=%.3f, model1_psnr=%.3f \r",
  472. i, #x,
  473. model1_time,
  474. math.sqrt(model1_mse / i), model1_psnr / i
  475. ))
  476. end
  477. end
  478. io.stdout:flush()
  479. end
  480. end
  481. if opt.save_info then
  482. local fp = io.open(path.join(opt.output_dir, "benchmark.txt"), "w")
  483. fp:write("options : " .. cjson.encode(opt) .. "\n")
  484. if baseline_psnr > 0 then
  485. fp:write(string.format("baseline: RMSE = %.3f, PSNR = %.3f\n",
  486. math.sqrt(baseline_mse / #x), baseline_psnr / #x))
  487. end
  488. if model1_psnr > 0 then
  489. fp:write(string.format("model1 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  490. math.sqrt(model1_mse / #x), model1_psnr / #x, model1_time))
  491. end
  492. if model2_psnr > 0 then
  493. fp:write(string.format("model2 : RMSE = %.3f, PSNR = %.3f, evaluation time = %.3f\n",
  494. math.sqrt(model2_mse / #x), model2_psnr / #x, model2_time))
  495. end
  496. fp:close()
  497. if detail_fp then
  498. detail_fp:close()
  499. end
  500. end
  501. io.stdout:write("\n")
  502. end
  503. local function load_data_from_dir(test_dir)
  504. local test_x = {}
  505. local files = dir.getfiles(test_dir, "*.*")
  506. for i = 1, #files do
  507. local name = path.basename(files[i])
  508. local e = path.extension(name)
  509. local base = name:sub(0, name:len() - e:len())
  510. local img = image_loader.load_float(files[i])
  511. if img then
  512. table.insert(test_x, {y = iproc.crop_mod4(img),
  513. basename = base})
  514. end
  515. if opt.show_progress then
  516. xlua.progress(i, #files)
  517. end
  518. end
  519. return test_x
  520. end
  521. local function load_data_from_file(test_file)
  522. local test_x = {}
  523. local files = utils.split(file.read(test_file), "\n")
  524. for i = 1, #files do
  525. local name = path.basename(files[i])
  526. local e = path.extension(name)
  527. local base = name:sub(0, name:len() - e:len())
  528. local img = image_loader.load_float(files[i])
  529. if img then
  530. table.insert(test_x, {y = iproc.crop_mod4(img),
  531. basename = base})
  532. end
  533. if opt.show_progress then
  534. xlua.progress(i, #files)
  535. end
  536. end
  537. return test_x
  538. end
  539. local function get_basename(f)
  540. local name = path.basename(f)
  541. local e = path.extension(name)
  542. local base = name:sub(0, name:len() - e:len())
  543. return base
  544. end
  545. local function load_user_data(y_dir, y_file, x_dir, x_file)
  546. local test = {}
  547. local y_files
  548. local x_files
  549. if y_file:len() > 0 then
  550. y_files = utils.split(file.read(y_file), "\n")
  551. else
  552. y_files = dir.getfiles(y_dir, "*.*")
  553. end
  554. if x_file:len() > 0 then
  555. x_files = utils.split(file.read(x_file), "\n")
  556. else
  557. x_files = dir.getfiles(x_dir, "*.*")
  558. end
  559. local basename_db = {}
  560. for i = 1, #y_files do
  561. basename_db[get_basename(y_files[i])] = {y = y_files[i]}
  562. end
  563. for i = 1, #x_files do
  564. local key = get_basename(x_files[i])
  565. if basename_db[key] then
  566. basename_db[key].x = x_files[i]
  567. else
  568. error(string.format("%s is not found in %s", key, y_dir))
  569. end
  570. end
  571. for i = 1, #y_files do
  572. local key = get_basename(y_files[i])
  573. local d = basename_db[key]
  574. if not (d.x and d.y) then
  575. error(string.format("%s is not found in %s", key, x_dir))
  576. end
  577. end
  578. for i = 1, #y_files do
  579. local key = get_basename(y_files[i])
  580. local x = image_loader.load_float(basename_db[key].x)
  581. local y = image_loader.load_float(basename_db[key].y)
  582. if x and y then
  583. table.insert(test, {y = y,
  584. x = x,
  585. basename = key})
  586. end
  587. if opt.show_progress then
  588. xlua.progress(i, #y_files)
  589. end
  590. end
  591. return test
  592. end
  593. function load_noise_scale_model(model_dir, noise_level, force_cudnn)
  594. local f = path.join(model_dir, string.format("noise%d_scale2.0x_model.t7", opt.noise_level))
  595. local s1, noise_scale = pcall(w2nn.load_model, f, force_cudnn)
  596. local model = {}
  597. if not s1 then
  598. f = path.join(model_dir, string.format("noise%d_model.t7", opt.noise_level))
  599. local noise
  600. s1, noise = pcall(w2nn.load_model, f, force_cudnn)
  601. if not s1 then
  602. model.noise_model = nil
  603. print(model_dir .. "'s noise model is not found. benchmark will use only scale model.")
  604. else
  605. model.noise_model = noise
  606. end
  607. f = path.join(model_dir, "scale2.0x_model.t7")
  608. local scale
  609. s1, scale = pcall(w2nn.load_model, f, force_cudnn)
  610. if not s1 then
  611. error(model_dir .. ": load error")
  612. return nil
  613. end
  614. model.scale_model = scale
  615. else
  616. model.noise_scale_model = noise_scale
  617. end
  618. return model
  619. end
  620. if opt.show_progress then
  621. print(opt)
  622. end
  623. if opt.method == "scale" or opt.method == "scale4" then
  624. local f1 = path.join(opt.model1_dir, "scale2.0x_model.t7")
  625. local f2 = path.join(opt.model2_dir, "scale2.0x_model.t7")
  626. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  627. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  628. if not s1 then
  629. error("Load error: " .. f1)
  630. end
  631. if not s2 then
  632. model2 = nil
  633. end
  634. local test_x
  635. if opt.file:len() > 0 then
  636. test_x = load_data_from_file(opt.file)
  637. else
  638. test_x = load_data_from_dir(opt.dir)
  639. end
  640. benchmark(opt, test_x, model1, model2)
  641. elseif opt.method == "noise" then
  642. local f1 = path.join(opt.model1_dir, string.format("noise%d_model.t7", opt.noise_level))
  643. local f2 = path.join(opt.model2_dir, string.format("noise%d_model.t7", opt.noise_level))
  644. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  645. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  646. if not s1 then
  647. error("Load error: " .. f1)
  648. end
  649. if not s2 then
  650. model2 = nil
  651. end
  652. local test_x
  653. if opt.file:len() > 0 then
  654. test_x = load_data_from_file(opt.file)
  655. else
  656. test_x = load_data_from_dir(opt.dir)
  657. end
  658. benchmark(opt, test_x, model1, model2)
  659. elseif opt.method == "noise_scale" then
  660. local model2 = nil
  661. local model1 = load_noise_scale_model(opt.model1_dir, opt.noise_level, opt.force_cudnn)
  662. if opt.model2_dir:len() > 0 then
  663. model2 = load_noise_scale_model(opt.model2_dir, opt.noise_level, opt.force_cudnn)
  664. end
  665. local test_x
  666. if opt.file:len() > 0 then
  667. test_x = load_data_from_file(opt.file)
  668. else
  669. test_x = load_data_from_dir(opt.dir)
  670. end
  671. benchmark(opt, test_x, model1, model2)
  672. elseif opt.method == "user" then
  673. local f1 = path.join(opt.model1_dir, string.format("%s_model.t7", opt.name))
  674. local f2 = path.join(opt.model2_dir, string.format("%s_model.t7", opt.name))
  675. local s1, model1 = pcall(w2nn.load_model, f1, opt.force_cudnn)
  676. local s2, model2 = pcall(w2nn.load_model, f2, opt.force_cudnn)
  677. if not s1 then
  678. error("Load error: " .. f1)
  679. end
  680. if not s2 then
  681. model2 = nil
  682. end
  683. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  684. benchmark(opt, test, model1, model2)
  685. elseif opt.method == "diff" then
  686. local test = load_user_data(opt.y_dir, opt.y_file, opt.x_dir, opt.x_file)
  687. benchmark(opt, test, nil, nil)
  688. end